clip_grad_norm_(max_norm, norm_type=2.0)
  Clips the gradient norm of all parameters. The norm is computed
  over all parameters' gradients as viewed as a single vector, and
  the gradients are modified in-place.

  Parameters:
     * **max_norm** (*float** or **int*) -- max norm of the
       gradients

     * **norm_type** (*float** or **int*) -- type of the used
       p-norm. Can be "'inf'" for infinity norm.

  Returns:
     Total norm of the parameters (viewed as a single vector).

  Return type:
     *Tensor*

  Note:

    If every FSDP instance uses "NO_SHARD", meaning that no
    gradients are sharded across ranks, then you may directly use
    "torch.nn.utils.clip_grad_norm_()".

  Note:

    If at least some FSDP instance uses a sharded strategy (i.e.
    one other than "NO_SHARD"), then you should use this method
    instead of "torch.nn.utils.clip_grad_norm_()" since this
