optimizer.zero_grad out of order

[jsta]

I believe the following chunk that appears throughout src/yoke/torch_training_utils.py:

optimizer.zero_grad(set_to_none=True)  # Possible speed-up
loss.mean().backward()
optimizer.step()

should instead be:

loss.mean().backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)  # Possible speed-up

as is, I think there is gradient leakage across epochs and higher than necessary usage overall:

with order corrected:

without order correction:

[jsta]

code is https://github.com/jsta/pytorch_memory/blob/main/create_profiler.py