Hi, please I have a similar problem, with the same kind of error message. I am new to this. everything works well with the forward propagation but the backprop keeps breaking after 2 epochs. Here is the code:

# Get cpu or gpu device for training.
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using {} device".format(device))



# Defining the Network model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        """this is the neural network 
        it can take a shape of 9 columns by any column
        the output layer is 2
        """
        self.model = nn.Sequential(
            nn.Linear(3*N,N-1),
            nn.Sigmoid(),
            nn.Linear(N-1,N-1),
            nn.Sigmoid(),
            nn.Linear(N-1,N-1),
            nn.Softmax(),
        )

   
    def forward(self, x, time_idx):
        
        #x is input data that you want to pass into the neural network
        # self.network_output is network output
        #t[:,time_idx+1] is the new clock time for the next time index (Eqn. 16) gotten from the weights (i.e., softmax o/p)
        self.network_output = self.model(x)
        
        for i in range(N):
            diff=t[:,time_idx]-t[i,time_idx]
            diff=diff[diff!=0]
            t[i,time_idx+1]=t[i,time_idx]+Time_res+ torch.sum(torch.dot(self.network_output[i,:], diff))
        self.newtime=t[:,time_idx+1]
        
        return self.network_output, self.newtime
    
    def custom_loss(self, t_outvec, time_idx):
        #t_outvec=self.newtime
        sumerror=torch.zeros((N,1), dtype=float)
        for i in range(N):
            ee=t_outvec-t_outvec[i]
            error= torch.sum(torch.square(ee))
            sumerror[i]=torch.log(torch.tensor(2+time_idx))*error
        self.loss=torch.sum(sumerror)
        return self.loss

modelmy=NeuralNetwork().to(device)
# Setting the input Matrix to the NN
Nodein=torch.zeros(N,3*N)
#Node=np.empty()
for b in range(N):
    Node=torch.cat([t[:,0], torch.Tensor(Distance_matrix[b,:]),torch.Tensor(RSSI_Matrix[b,:])])
    Nodein[b,:]=Node

torch.autograd.set_detect_anomaly(True)
#loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(modelmy.parameters(), lr=0.01)
steps = 20
for i in range(steps):
    output = modelmy.forward(Nodein,i)
    loss=modelmy.custom_loss(output[1],i)
    print(loss)
    #loss = loss_fn(y_true, output)
    #print(f"loss = {loss}")
    optimizer.zero_grad()
    loss.backward(retain_graph=True)
    optimizer.step()

I will greatly appreciate your help