code: def get_cur_mem(rank, result, prefix):
      """Collect memory allocated values in a result dict in MB"""
      result[prefix] = round(
          torch.cuda.memory_allocated() / 1024 / 1024
      )
  
imports: import contextlib
  import pytest
  import torch
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  import torch.optim
  from fairscale.fair_dev.testing.testing import dist_init, dump_all_tensors, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.internal import torch_version
  from fairscale.internal.parallel import get_process_group_cached
  from fairscale.nn import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import auto_wrap_bn
label_desc: Use image node_id 2 for calling a global function with example usage: get_cur_mem(rank, result, prefix) without return types