code: def _test_memory(self, config, rank, group, volatile=False):
      model = self.get_wrapped_model(
          group, cuda_first=False, config=config
      )
      self._train_for_several_steps(
          model, 1, autocast=model.mixed_precision
      )
  
      mems = [get_cuda_mem()]
  
      with model.summon_full_params(volatile=volatile):
          mems.append(get_cuda_mem())
          assert mems[1] >= mems[0]
  
          state_dict = model.state_dict()
          mems.append(get_cuda_mem())
          assert mems[2] >= mems[1]
  
      mems.append(get_cuda_mem())
      assert mems[3] <= mems[2]
  
      del state_dict
      mems.append(get_cuda_mem())
      # Any value other than `==` indicates a memory leak. If mems[4] >
      # mems[0], that indicates we're not cleaning up params properly in
      # summon_full_params. If mems[4] < mems[0], that indicates there's a
      # memory leak in _train_for_several_steps.
      assert (
          mems[4] == mems[0]
      ), f"memory leak detected, {mems[4]} != {mems[0]}"
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_summon_full_params.py
imports: import functools
  import gc
  import unittest
  from parameterized import parameterized
  import pytest
  import torch
  from fairscale.internal.version import torch_version
  from .test_fsdp import CONFIG_OPTIONS, DistributedTest, rename_test, spawn_and_init
label_desc: Use image node_id 3 for calling the TestMemory obj's underlying member method code with example usage: obj._test_memory(config, rank, group, volatile) without return types