class_docstr: Tests ``fully_shard`` initialization.
class_docstr_tok: ["Tests","``","fully_shard","``","initialization","."]
code: def _test_policy(self, policy: Optional[_Policy]):
      use_nested_sequential_model = "Sequential" in getattr(
          policy, "_module_classes_str", ""
      )
      local_model = (
          NestedSequentialModel(torch.device("cuda"))
          if use_nested_sequential_model
          else CompositeParamModel(torch.device("cuda"))
      )
      fsdp_wrapped_model = FSDP(
          copy.deepcopy(local_model),
          auto_wrap_policy=policy,
          use_orig_params=True,
      )
      composable_module = copy.deepcopy(local_model)
      fully_shard(
          composable_module,
          policy=policy,
      )
      self._test_fully_shard_construction(
          local_model,
          fsdp_wrapped_model,
          composable_module,
      )
  
imports: import copy
  import sys
  from typing import Optional
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._composable import fully_shard
  from torch.distributed.fsdp import BackwardPrefetch, FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _is_fsdp_flattened, clean_tensor_name
  from torch.distributed.fsdp.wrap import _Policy, CustomPolicy, ModuleWrapPolicy
  from torch.testing._internal.common_dist_composable import CompositeParamModel, FakeSequential, NestedSequentialModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 3 for calling the TestInitialization obj's underlying member method code with example usage: obj._test_policy(policy) without return types