Pytorch leaf Variable that requires grad is being used in an in-place operation in a experimental loss function

Hi everyone I'm new to all this pytorch stuff. I am trying for experimental reasons, to define the following loss function that would measure, so to speak, the convergence of the model to a uniform data distribution, which would mesure, so to speak, in which quantile the real data falls with respect to K simulations. That is, the model generates K simulations and the real data will be greater than j of them. If the model is well trained this should converge to the uniform distribution.

class CustomLoss(nn.Module):
    def __init__(self, K=5):
        super(CustomLoss, self).__init__()
        self.K = K
        self.a_k = torch.zeros(K+1).requires_grad_(True)
        self.n = 0

    def forward(self, ys_hat, ys_training):
        for k in range(self.K+1):
            for y_training in ys_training:
                self.a_k[k] += test_m(phy(ys_hat, y_training), k)
        self.n+=1
        return (self.a_k - (self.n/self.K)).pow(2).sum()

However when I train this, I get the following error:

 a view of a leaf Variable that requires grad is being used in an in-place operation.

However, I don't know how I can avoid this as the self.a_k will always be leaf and I want to save them for reuse in later optimizer calls.

source https://stackoverflow.com/questions/76213380/pytorch-leaf-variable-that-requires-grad-is-being-used-in-an-in-place-operation