I’ve implemented some custom functions in CUDA. They pass gradcheck individually, but when I compose them, they fail. Is this expected behavior? I can’t imagine it would be. The structure of my function is:
import torch
from my_src_code import my_backend1, my_backend2
class MyFunction(torch.autograd.Function):
@staticmethod
def forward(self, input):
self.save_for_backward(input)
out = my_backend1.forward(input)
return my_backend2.forward(out)
@staticmethod
def backward(self, grad_output):
input = self.saved_tensors
grad_output = my_backend2.backward(grad_output)
grad_input = my_backend1.backward(
grad_output,
input)
return grad_input
class MyLayer(nn.Module):
def __init__(self):
super(MyLayer, self).__init__()
# ...
def forward(self, x):
return MyFunction.apply(x)
Both my_backend1
and my_backend2
pass gradcheck on their own.