aberk:

self.weights = F.normalize(self.weights, p=2, dim=0)


How about change this line into
self.weights = nn.Parameter(F.normalize(self.weights, p=2, dim=0))

According to the partial snippet, the retain_graph error may be caused by there is some computation in global and cleared in the first .backward that when you backward twice it cannot get this node to compute the gradient.