This is the code of the training process
optimizer = optim.Adam(net.parameters(), lr = 0.001)
criterion = torch.nn.MSELoss() 
EPOCHS = 500000

lloss = []

for epoch in range(EPOCHS):
    avg_train_loss = 0
    for trainbatch in train_loader:
        X,y = trainbatch
        net.zero_grad()
        output = net(X.float())
        loss = criterion(output[:, -1], y)
        loss.backward()
        optimizer.step()
        avg_train_loss += loss.item() / len(train_loader)
    lloss.append(avg_train_loss)
        

There is no no_grad() there, and I think it has nothing to do with the problem I presented