code: def test_freezing_weights(
      self,
      with_nested_trunk,
      freezing_method,
      freeze_after_wrap_fsdp,
      disable_autograd,
      forward_prefetch,
  ):
      # DDP
      ddp_state = self._dist_train(
          with_nested_trunk,
          freezing_method,
          freeze_after_wrap_fsdp,
          with_fsdp=False,
          disable_autograd=disable_autograd,
          forward_prefetch=False,  # does not apply to DDP
      )
  
      # FSDP
      fsdp_state = self._dist_train(
          with_nested_trunk,
          freezing_method,
          freeze_after_wrap_fsdp,
          with_fsdp=True,
          disable_autograd=disable_autograd,
          forward_prefetch=forward_prefetch,
      )
  
      self.assertEqual(
          ddp_state,
          fsdp_state,
          exact_device=True,
          msg="FullyShardedDataParallel states didn't match PyTorch DDP states",
      )
  
      if freezing_method == FreezingMethod.RequiresGrad:
          for ddp_param, fsdp_param in zip(ddp_state, fsdp_state):
              self.assertEqual(
                  ddp_param.requires_grad, fsdp_param.requires_grad
              )
  
imports: import contextlib
  import sys
  from enum import Enum
  import torch
  import torch.nn
  import torch.optim
  from torch import distributed
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest, get_full_params
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 3 for calling the TestFreezingWeights obj's underlying member method code with example usage: obj.test_freezing_weights(with_nested_trunk, freezing_method, freeze_after_wrap_fsdp, disable_autograd, forward_prefetch) without return types