code: def test_raises_on_mismatch(self):
      torch.manual_seed(self.seed + self.rank)
  
      # Identical models with data-parallel and ZeRO
      num_layers = 11
      layer_size = 7
      ref_model, ref_optim, dist_model, dist_optim = make_models(
          num_layers,
          layer_size,
      )
  
      # Only perform training step with distributed model
      dist_optim.zero_grad()
      x = torch.rand(3, layer_size) - 0.5
      x = x.to(dtype=torch.float32, device="cuda")
      dy = torch.rand_like(x) - 0.5
      y = dist_model(x)
      y.backward(dy)
      dist_optim.step()
  
      # Check that parameters do not match
      for ref_param, dist_param in zip(
          ref_model.parameters(), dist_model.parameters()
      ):
          self.assertRaises(
              AssertionError,
              torch.testing.assert_close,
              dist_param,
              ref_param,
          )
  
imports: from contextlib import contextmanager
  import io
  from typing import Callable, Optional, Tuple
  import unittest
  import warnings
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 16 for calling the TestDistributedFusedAdam obj's underlying member method code with example usage: obj.test_raises_on_mismatch() without return types