code: def get_tensors_by_size_bucket():
      size_buckets = defaultdict(int)
      for obj in gc.get_objects():
          if not isinstance(obj, torch.Tensor):
              continue
          if obj.device.type == "cuda":
              size_buckets[(*obj.size(),) + (obj.element_size(),)] += 1
  
      return size_buckets
  
imports: from collections import defaultdict
  import gc
  import logging
  import math
  import time
  import torch
  import torch.distributed
  from torch.distributed import rpc
  from torch.nn.parallel import DistributedDataParallel
  import utils
  from benchmarks.golden_configs.lm_wikitext2 import Pipe
  from fairscale.fair_dev.testing.testing import dist_init
  from fairscale.nn import Pipe
  from fairscale.nn.model_parallel import initialize_model_parallel
label_desc: Use image node_id 1 for calling a global function with example usage: get_tensors_by_size_bucket() and returns: size_buckets