docs: Fix order of zero_grad

[ValerianRey]

The idea is to free the .grad fields as early as we can in our usage examples, so that memory consumption is reduced for users that would follow these examples. See https://discuss.pytorch.org/t/order-of-backward-step-and-zero-grad/119023.

It also changes a test to use optimizer.zero_grad instead of model.zero_grad. Both are equivalent but I'd rather always use the former so that it's easier to find and replace.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/torchjd/autogram/_engine.py	100.00% <ø> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.