code: def allocate_tensor(self, size):
      free_before = self.total_free
  
      assert (
          size <= self.total_free
      ), "Not enough memory in buffer. Allocation failed"
      if self.largest_contiguous < size:
          print_rank_0(
              "Needs defragmentation to allocate. Before Defragmentation:"
          )
          self.print_allocation(resolution=100)
          self._defragment_memory()
          # set the param data to the new tensor buffer locations
          self._reset_param_data()
          print_rank_0("After defragmentation:")
          self.print_allocation(resolution=100)
  
      self.total_free = self.total_free - size
  
      allocated = self.total_size - self.total_free
      if allocated > self.max_allocated:
          self.max_allocated = allocated
  
      tensor_address = self._get_new_tensor_address(size)
  
      ret_tensor = self._get_new_tensor(tensor_address, size)
      print_rank_0(
          f"Free before allocation {free_before}. Allocating {size}. Free after allocation {self.total_free}. Max allocated {self.max_allocated}"
      )
      assert (
          self.total_free + size == free_before
      ), "Allocation bookkeeping error"
  
      return ret_tensor
  
imports: import torch
  from deepspeed import comm
label_desc: Use image node_id 2 for calling the ContiguousMemoryAllocator obj's underlying member method code with example usage: obj.allocate_tensor(size) and returns: ret_tensor