toho:

output = net(A, B) loss = loss_fn(output.squeeze(), labels) optimizer.zero_grad() loss.backward() optimizer.step()


My mistake was that I provided the data, computed the loss and then zero-ed the grads, did back-prop and step()… If you place the ‘provide data’ call correctly, this is fine :).