code: def _test_cudnn_gbn(
      self,
      num_layers: int,
      shape: typing.List[int],
      *,
      memory_format: torch.memory_format = torch.channels_last,
  ) -> None:
      global_shape = copy.deepcopy(shape)
      global_shape[0] = self.world_size
  
      device = torch.device("cuda", self.rank)
      cudnn_gbn_model = BNModel(
          num_features=shape[1],
          num_layers=num_layers,
      ).to(device=device, memory_format=memory_format)
      ref_model = BNModelRef(
          num_features=shape[1],
          num_layers=num_layers,
      ).to(device=device, memory_format=memory_format)
  
      input, weight, bias, grad_out = get_rand_tensors(
          global_shape, device
      )
      with torch.no_grad():
          ref_model.fwd[0].weight.copy_(weight)
          ref_model.fwd[0].bias.copy_(bias)
          cudnn_gbn_model.fwd[0].weight.copy_(weight)
          cudnn_gbn_model.fwd[0].bias.copy_(bias)
  
          ref_input = input.clone().detach().requires_grad_()
          input = (
              input[self.rank : self.rank + 1, ...]
              .clone()
              .detach()
              .requires_grad_()
          )
  
          ref_grad_out = grad_out.half().clone().detach()
          grad_out = (
              grad_out[self.rank : self.rank + 1, ...]
              .half()
              .clone()
              .detach()
          )
  
      with torch.cuda.amp.autocast():
          out = cudnn_gbn_model(input)
          ref_out = ref_model(ref_input.half())
      out.backward(grad_out)
      ref_out.backward(ref_grad_out)
  
      kwargs = {"rtol": 3.5e-3, "atol": 3e-2, "msg": f"shape: {shape}"}
  
      torch.testing.assert_close(
          ref_out[self.rank : self.rank + 1], out, **kwargs
      )
      torch.testing.assert_close(
          ref_input.grad[self.rank : self.rank + 1],
          input.grad,
          **kwargs,
      )
      # compensating the averaging over processes done by DDP
      # in order to produce mathematically equivalent result
      # https://github.com/NVIDIA/apex/issues/134#issuecomment-458307368
      torch.testing.assert_close(
          ref_model.fwd[0].weight.grad / self.world_size,
          cudnn_gbn_model.fwd[0].weight.grad,
          **kwargs,
      )
      torch.testing.assert_close(
          ref_model.fwd[0].bias.grad / self.world_size,
          cudnn_gbn_model.fwd[0].bias.grad,
          **kwargs,
      )
  
imports: import copy
  import typing
  import unittest
  import torch
  import torch.nn
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 3 for calling the TestCudnnGBN obj's underlying member method code with example usage: obj._test_cudnn_gbn(num_layers, shape) without return types