code: def _freeze_distributed_worker(
      gpu_id,
      world_size,
      tempfile_name,
      unused,
  ):
      torch.cuda.set_device(gpu_id)
  
      rank = gpu_id
      result = dist_init(rank, world_size, tempfile_name, unused)
      assert result, "Dist init failed"
  
      torch.manual_seed(0)
      torch.backends.cudnn.deterministic = True
      batch = torch.randn(size=(2, 3, 224, 224)).cuda()
  
      # The use case for this test is where the weights in the submodule
      # are not frozen but the leftover weights or those contained by the
      # root module are frozen. Refer to issue #758 for a real world example.
      model = FreezeModel()
      model = model.cuda()
  
      for param in model.head.parameters():
          param.requires_grad = False
  
      model = FSDP(model)
  
      if gpu_id == 0:
          print(model)
  
      target = torch.tensor([0, 1], dtype=torch.long).cuda()
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
  
      for iteration in range(3):
          out = model(batch)
          fake_loss = criterion(out, target)
          print("Loss", iteration, ":", fake_loss.item())
          optimizer.zero_grad()
          fake_loss.backward()
          optimizer.step()
  
      teardown()
  
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 1 for calling a global function with example usage: _freeze_distributed_worker(gpu_id, world_size, tempfile_name, unused) without return types