# https://github.com/pytorch/pytorch/blob/6dc445e1a84dc5d093d640de54f038f021d13227/torch/nn/modules/module.py#L192-L208
def _apply(self, fn):
    ...
    for param in self._parameters.values():
        if param is not None:
            # Tensors stored in modules are graph leaves, and we don't
            # want to create copy nodes, so we have to unpack the data.
            param.data = fn(param.data)  # NOTE: this doesn't allow changing `param.data`'s TensorImpl type
            if param._grad is not None:
                param._grad.data = fn(param._grad.data)  # NOTE: this doesn't allow changing `param._grad.data`'s TensorImpl type
   ...

def _apply(self, fn):
    ...
    for key, param in self._parameters.items():
        if param is not None:
            # Tensors stored in modules are graph leaves, and we don't
            # want to create copy nodes, so we have to unpack the data.
            param_applied = fn(param)
            if param._is_same_impl_type(param_applied):
                param.data = param_applied
            else:  # If we have to change TensorImpl type...
                with torch.no_grad():
                    # We use `requires_grad_()` here, to make sure the new `param` still
                    # has the same `requires_grad` value as the old `param`. An alternative is
                    # to not use `with torch.no_grad():`, but that would cause the following operation
                    # to create a `CopyBackwards` gradient function which is not what we wanted.
                    self._parameters[key] = param_applied.requires_grad_(param.requires_grad)
            if param._grad is not None:
                grad_applied = fn(param._grad)
                if param._grad._is_same_impl_type(grad_applied):
                    param._grad.data = grad_applied
                else:  # If we have to change TensorImpl type...
                    self._parameters[key]._grad = grad_applied
   ...