code: def _reshard_output(
      module: torch.nn.Module, resharding_spec: ShardingSpec
  ) -> torch.nn.Module:
      """
      Hook a module with output resharding in the forward pass according
      to the given ``resharding_spec``.
  
      Args:
          module (:class:`torch.nn.Module`): Module whose output needs to be resharded.
          resharding_spec (:class:`torch.distributed._shard.sharding_spec.ShardingSpec`):
              The specification describing how the output of the module will be resharded.
  
      Returns:
          A :class:`torch.nn.Module` object with reshard API hooked.
      """
  
      def hook_func(_module, _input, output):
          if isinstance(output, ShardedTensor):
              return output.reshard(resharding_spec)
          return output
  
      module.register_forward_hook(hook_func)
      return module
  
file_path: pytorch/torch/distributed/_shard/api.py
imports: from contextlib import contextmanager
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed import distributed_c10d
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from .sharding_spec import ShardingSpec, ChunkShardingSpec
  from .sharding_plan import ShardingPlan
  from .sharder import Sharder
label_desc: Use image node_id 5 for calling a global function with example usage: _reshard_output(module, resharding_spec) and returns: module, output, output