Small docstring fix (#4431)
* small docstring fix
* default is no longer none in zero_grad
* typo
* max_norm not in method signature
* Revert "default is no longer none in zero_grad"
This reverts commit cb126bae0820392f3afcd42220b0639a5fbdd4b3.
---------
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
