code: def shard_fn(name, module, device_mesh):
      if isinstance(module, nn.Linear):
          for name, param in module.named_parameters():
              dist_param = torch.nn.Parameter(
                  distribute_tensor(param, device_mesh, [Shard(0)])
              )
              # make sure partial sum get cleared after backward()
              dist_param.register_hook(
                  lambda grad: grad.redistribute(placements=[Shard(0)])
              )
              module.register_parameter(name, dist_param)
  
file_path: pytorch/test/distributed/_tensor/test_optimizers.py
imports: from copy import deepcopy
  import torch
  import torch.nn
  from torch.distributed._tensor import DeviceMesh, distribute_module, distribute_tensor, DTensor, Replicate, Shard
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, MLPModule, with_comms
label_desc: Use image node_id 1 for calling a global function with example usage: shard_fn(name, module, device_mesh) without return types