@vmoens As currently implemented, the new hook's signature is different from other hooks: besides accepting (and returning) a TensorDict instance, it also accepts the clip norm parameters and the index i in the trainer._optimizer_ops list (for distinguishing each hook's grad_norm_{i} key added to the returned TensorDict).

If we want the optimizer hook to be more similar to other hooks, we could:

1. Add clip_norm and clip_norm_grad constructor parameters to the OptimizerHook.__init__. Then, when creating the "default" optimizer hook from its optimizer argument, Trainer.__init__ will just pass its own clip_norm and clip_norm_grad arguments as is to the OptimizerHook.__init__. Bonus: users will be able to set different gradient clipping configuration for each hook separately.

2. Return grad_norm value from the OptimizerHook.__call__ instead of the updated TensorDict parameter. Then in Trainer._optimizer_hook we could set the grad_norm in the losses_td with the respective index, without the need to pass the index to the hook call.

3. Move detaching the losses_td outside of Trainer._optimizer_hook: to be done by the calling Trainer.optim_steps right after the Trainer._optimizer_hook returns.

Importantly, even after the above modifications, Trainer._optimizer_hook will be different from other Trainer._*_hook functions, as it will need to set grad_norm in its TensorDict argument besides the usual for-loop over the hook instances.

Please let me know how you'd like to proceed. Thanks!

Codecov Report

Merging #716 (860e922) into main (3105819) will increase coverage by 0.12%.
The diff coverage is 97.87%.

@@            Coverage Diff             @@
##             main     #716      +/-   ##
==========================================
+ Coverage   88.77%   88.89%   +0.12%     
==========================================
  Files         122      122              
  Lines       21151    21273     +122     
==========================================
+ Hits        18776    18910     +134     
+ Misses       2375     2363      -12     

Flags with carried forward coverage won't be shown. Click here to find out more.

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

Regarding this

Importantly, even after the above modifications, Trainer.optimizer_hook will be different from other Trainer.*_hook functions, as it will need to set grad_norm in its TensorDict argument besides the usual for-loop over the hook instances.

It's ok, it's backward compatible and it suits most needs.
Thanks again for this!