code: def test_freezing_weights(temp_files, nested_trunk):
      with_nested_trunk = nested_trunk == "nested_trunk"
  
      world_size = 2
      # DDP
      with_fsdp = False
      freezing_method = FreezingMethod.RequiresGrad
      mp.spawn(
          _distributed_worker,
          (
              world_size,
              with_fsdp,
              with_nested_trunk,
              freezing_method,
              True,
          )
          + temp_files[0:3]
          + (None,),
          nprocs=world_size,
      )
      # FSDP, case 1 and 2.
      with_fsdp = True
      expected_state = torch.load(temp_files[2])
      temp_file_idx = 3
      for freezing_method, freeze_after_wrap_fsdp in product(
          [FreezingMethod.RequiresGrad, FreezingMethod.GradToNone],
          [True, False],
      ):
          print(f"Testing FSDP with freezing method {freezing_method}")
          mp.spawn(
              _distributed_worker,
              (
                  world_size,
                  with_fsdp,
                  with_nested_trunk,
                  freezing_method,
                  freeze_after_wrap_fsdp,
              )
              + temp_files[temp_file_idx : temp_file_idx + 3]
              + (expected_state,),
              nprocs=world_size,
          )
          temp_file_idx += 3
  
imports: from enum import Enum
  from itertools import product
  import tempfile
  import pytest
  import torch
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, objects_are_equal, rmf, skip_if_single_gpu, teardown
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 6 for calling a global function with example usage: test_freezing_weights(temp_files, nested_trunk) without return types