Get import error when compile a third op by using CppExtension

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)
  ;
}

@Shan19900305 did you manage to solve this?

On my end adding

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
}

solves the issue, but then on python side one still needs import myCopyOp so as to use torch.ops.my_copy_function.MyCopyFunction.