class_docstr: Tests ``fully_shard`` initialization.
class_docstr_tok: ["Tests","``","fully_shard","``","initialization","."]
docstr: Tests passing a ``policy`` for pseudo-auto-wrapping.
docstr_tok: ["Tests","passing","a","``","policy","``","for","pseudo-auto-wrapping","."]
code: def test_policy(self):
      """Tests passing a ``policy`` for pseudo-auto-wrapping."""
  
      def lambda_fn(module: nn.Module):
          if isinstance(module, nn.Sequential):
              return True
          elif isinstance(module, FakeSequential):
              return {
                  "backward_prefetch": BackwardPrefetch.BACKWARD_POST
              }
          return False
  
      self.run_subtests(
          {
              "policy": [
                  None,
                  ModuleWrapPolicy({UnitModule}),
                  ModuleWrapPolicy({nn.Sequential}),
                  CustomPolicy(lambda_fn),
              ],
          },
          self._test_policy,
      )
  
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 2 for calling the TestInitialization obj's underlying member method code with example usage: obj.test_policy() and returns: False, True, dict