I would like to minimize the following cost function
class SelfSupervLoss(nn.Module):
    def __init__(self):
        super(SelfSupervLoss, self).__init__()
        
    def forward(self, x, edge_index):
        src, dest = edge_index
        return (x[src] * x[dest]).mean()

criterion = SelfSupervLoss()

and I want to use SGD for this purpose
from torch.optim import SGD

optimizer = SGD(model.parameters(), lr = 0.01)

The only parameters of my model are:
self.hv = torch.nn.Parameter(data.x[data.edge_index[0]]) # argument is a Tensor with the shape [472, 6]
self.hu = torch.nn.Parameter(data.x[data.edge_index[1]]) # argument is a Tensor with the shape [472, 6]

which I declare in the subclass of torch.nn.Module (If needed I can share the whole model)
Training loop is given as follows:
def train():
    optimizer.zero_grad()
    out = model(data.x, data.edge_index, deg)
    loss = criterion(out, data.edge_index)
    loss.backward(retain_graph = True)
    optimizer.step()
    return loss

But the loss remains on the same level over 200+ epochs
 Epoch: 0, Loss: 0.1043
 Epoch: 1, Loss: 0.1087
 Epoch: 2, Loss: 0.0914
 Epoch: 3, Loss: 0.1007
 Epoch: 4, Loss: 0.0994
...

I do not understand what I can be doing wrong. Maybe the declaration of the parameters? I am not sure… Does anybody know?
Thank you!