This is proposed as a simple fix for #6595. It follows the torch-like practice of wrapping functions for pre and post hooks (as also followed by many other repos) with arg and kwargs for argument forwarding. See example:

https://github.com/pytorch/pytorch/blob/4cfa06f706811e2c3b70289b222461822537bd93/torch/nn/modules/module.py#L1498

This seems like it would solve the problem. WDYT?

Description

A few sentences describing the changes proposed in this pull request.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.