I am just starting to learn Pytorch, and trying to build a custom loss function
class distance_loss(torch.nn.Module):
  def __init__(self, in_features = 10, out_features = 1):
    super(distance_loss, self).__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.center = torch.nn.Parameter(torch.randn(in_features))
    self.center.requires_grad = True

  def forward(self, input):
    return torch.norm(input - self.center, 2)

Then I want to use it in some model:
model = SomeModel()

loss = distance_loss()

scores = model(torch.rand(1,3,32,32))

optimizer = torch.optim.SGD(model.parameters(), lr = 0.001 )

optimizer.zero_grad()

loss.backward()

optimizer.step()

But this returns error: “distance_loss” object has no attribute ‘backward’. It is my understanding that backward method is automatically implemented in nn.module. What goes wrong here? Also do I somehow need to append the learnable parameters of distance_loss() to model.parameters() for optimizer to know?