Hi Jack!
And also be sure to zero out your gradients after taking the optimization
step.  (Otherwise each loss.backward() call will keep accumulating
into the gradient.)  E.g., something like this:
self.fully_connected_layer_1_weight.grad.zero_()

Best.
K. Frank