code: def test_norm(device, norm_type, mixed_precision):
      """Test checkpoint_wrapper with different norm layers."""
      if device == "cuda" and not torch.cuda.is_available():
          pytest.skip("Skip due to lack of GPU")
  
      # Get input, ref, checkpoint models and make them equal.
      in_data = torch.rand(2, 2, 3, 3).to(device)
      m_ref = get_model(norm_type, False, mixed_precision).to(device)
      m_cpt = get_model(norm_type, True, mixed_precision).to(device)
      m_cpt.load_state_dict(m_ref.state_dict())
  
      if torch_version() >= (1, 6, 0):
          # This assert fails on 1.5.1.
          assert objects_are_equal(
              m_ref.state_dict(), m_cpt.state_dict()
          )
  
      if mixed_precision != "fp32":
          in_data = in_data.half()
  
      # Needed due to checkpointing.
      in_data.requires_grad = True
      for model in (m_ref, m_cpt):
          optim = SGD(model.parameters(), lr=0.1)
          if device == "cpu" and mixed_precision != "fp32":
              # Got: RuntimeError: "batch_norm"/"layer_norm" not implemented for 'Half'.
              with pytest.raises(RuntimeError):
                  out = model(in_data)
              return
          else:
              # Everything else work.
              out = model(in_data)
          out.sum().backward()
          optim.step()
  
      if torch_version() >= (1, 6, 0):
          assert objects_are_equal(
              m_ref.state_dict(), m_cpt.state_dict()
          )
  
file_path: fairscale/tests/nn/checkpoint/test_checkpoint_activations_norm.py
imports: import pytest
  import torch
  from torch.nn import BatchNorm2d, LayerNorm, Linear, Sequential
  from torch.optim import SGD
  from fairscale.fair_dev.testing.testing import objects_are_equal
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
label_desc: Use image node_id 2 for calling a global function with example usage: test_norm(device, norm_type, mixed_precision) without return types