I modify a custom autograd function like this
I thought this error is caused by multiple return values from function A
but I need min scale to make other codes run properly.
class A(torch.autograd.Function):
@staticmethod
def forward(ctx, input, qparams=None):
…
return output, min, scale
@staticmethod
def backward(ctx, grad_output):
grad_input = grad_output
return grad_input, None
class B(nn.Module):
def init(self):
…
def forward(self, input, num_bits, qparams=None):
…
output, min, scale = A().apply(input, qparams)
return output