pred = model(x)[0]
        flat_pred = torch.reshape(pred, (pred.size(0) * pred.size(1), pred.size(2)))
        #flattening for batch loss computation
        loss =  mse(torch.argmax(flat_pred, dim = 1).double(), y.flatten().double()) #+  nll(flat_pred, y.flatten())
        total_loss.append(loss)

If I print the requires_grad, I see that it is true for the flat_pred tensor, but torch.argmax(flat_pred, dim = 1).double().requires_grad is False