Get import error when compile a third class op by using CppExtension.
A op class is inherited from CustomClassHolder, and using macro TORCH_LIBRARY to bind to pytorch. But after compiled by CppExtension, this module can’t import to python side.
Anyway to fix this error?
Error message:
python -c "import myCopyOp"
Traceback (most recent call last):
File "<string>", line 1, in <module>
ImportError: dynamic module does not define module export function (PyInit_myCopyOp)
setup.py
from setuptools import setup, find_packages
from torch.utils import cpp_extension
setup(name="myCopyOp",
packages=find_packages(),
ext_modules=[cpp_extension.CppExtension('myCopyOp',
sources=['self_ops.cpp'],
)],
cmdclass={'build_ext': cpp_extension.BuildExtension})
C++ code:
#include "torch/torch.h"
#include "torch/custom_class.h"
class MyCopyFunction : public torch::CustomClassHolder {
public:
MyCopyFunction(int64_t size): size_(size) {
weight_ = torch::randn({size}, at::kCPU);
}
at::Tensor forward(const at::Tensor& self) {
at::Tensor dst = torch::mul(self, weight_);
return dst;
}
private:
at::Tensor weight_;
const int64_t size_;
};
// PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// m.def("forward", &MyCopyFunction::forward, "LLTM forward");
// }
TORCH_LIBRARY(my_copy_function, m) {
m.class_<MyCopyFunction>("MyCopyFunction")
.def(torch::init<int64_t>())
.def("forward", &MyCopyFunction::forward)
;
}