code: def test_replicate_ignore_module(self):
      self._init_pg()
      torch.cuda.set_device(self.rank)
      # Seed ensures diff input and thus different local grads across ranks.
      torch.manual_seed(self.rank)
      torch.cuda.manual_seed(self.rank)
      model = Net().cuda()
      replicate(model, ignored_modules=[model.fc1])
      # CPU input ensures that replicate can move input to GPU as DDP does.
      inp = torch.randn(5, 2, device="cuda") * (self.rank + 1)
      out = model(inp) * 10
      out.sum().backward()
      # FC1 grads should not be synchronized, FC2 and 3 should be.
      fc1_grad = model.fc1.weight.grad
      tensor_list = [
          torch.zeros_like(fc1_grad)
          for _ in range(dist.get_world_size())
      ]
      dist.all_gather(tensor_list, fc1_grad)
      grad, rest = tensor_list[0], tensor_list[1:]
      for g in rest:
          self.assertNotEqual(grad, g)
  
      for dp_grad in [model.fc2.weight.grad, model.fc3.weight.grad]:
          tensor_list = [
              torch.zeros_like(dp_grad)
              for _ in range(dist.get_world_size())
          ]
          dist.all_gather(tensor_list, dp_grad)
          grad, rest = tensor_list[0], tensor_list[1:]
          for g in rest:
              self.assertEqual(grad, g)
  
imports: import os
  from copy import deepcopy
  import torch
  import torch.distributed
  import torch.nn.functional
  from torch import nn
  from torch.distributed._composable.replicate import replicate
  from torch.testing._internal.common_distributed import MultiProcessTestCase, skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests
label_desc: Use image node_id 8 for calling the ReplicateTest obj's underlying member method code with example usage: obj.test_replicate_ignore_module() without return types