The code looks generally fine.
One check you could additionally do is to check for valid gradients in all modules.
E.g.:
...
optimizer.step()
print(net.some_layer.weight.grad)
# or 
for name, param in net.named_parameters():
    print(name, param.grad)
# the same for classifier