sglang_v0.5.2/pytorch_2.8.0/test/custom_operator/op.h

25 lines
530 B
C++

#include <torch/script.h>
#include <cstddef>
#include <vector>
#include <string>
// clang-format off
# if defined(_WIN32)
# if defined(custom_ops_EXPORTS)
# define CUSTOM_OP_API __declspec(dllexport)
# else
# define CUSTOM_OP_API __declspec(dllimport)
# endif
# else
# define CUSTOM_OP_API
# endif
// clang-format on
CUSTOM_OP_API torch::List<torch::Tensor> custom_op(
torch::Tensor tensor,
double scalar,
int64_t repeat);
CUSTOM_OP_API int64_t custom_op2(std::string s1, std::string s2);