Thank you for your suggestion. Although I don’t why it did work in my case, I checked the keyword “register” and found I should add the below code at the end. It works.

TORCH_LIBRARY(my_ops, m) {
m.def(“SpikeFunction”, &SpikeFunction);
}

But I have a following question; the returned tensor’s requires_grad is false and I cannot backward it properly. Initially, I thought it may be due to the logical format then I tried a square operation but the issue emerged.




a=torch.rand(5,requires_grad=True)
b=torch.ops.my_ops.SpikeFunction(a,torch.tensor(0.3))
b[0]
tensor([True, True, True, True, True])




b[0].requires_grad
False




By the way, may I ask how to cast tensor from logic to float in c++? I checked the doc but I don’t know how to use it. I tried torch::greater(v_scaled, 0.)._cast_Float() ,   torch::greater(v_scaled, 0.) .at::_cast_Float() ,  at ::_cast_Float( torch::greater(v_scaled, 0.)) and many other patterns, all do not work.
Thank you!