Hi -
I’m trying to implement a custom loss function. Not differentiable, requires I implement gradients. The interpreter doesn’t like the current foo model:
class my_loss_fn(Function):

    def __init__(self, scalar=None):
        super(my_loss_fn, self).__init__()
        self.scalar = scalar
    
    @staticmethod
    def forward(self, i):
        out = (i - i).pow(2).sum()
        self.mark_non_differentiable(out)
        return out
    
    @staticmethod
    def backward(self, grad_output):
        grad_input = grad_output.clone()
        grad_input = 2.0*(grad_output)
        return grad_input, None

model = nn.Linear(2, 2)
x = torch.randn(1, 2)
target = torch.randn(1, 2)
output = model(x)
loss = my_loss_fn.apply(output)
loss.backward()
print(model.weight.grad)

I get error: RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
I’ve tried throwing requires_grad = True in some places, per some suggestions in other posts, with no luck.
What am I missing?
I don’t care if the model does anything meaningful, for now.