class_docstr: Persistent buffer for main gradients that remains allocated between training iterations
code: def __init__(self, numel, chunk_size_mb):
      if not HAVE_APEX:
          raise ImportError(
              "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      if not HAVE_MEGATRON_CORE:
          raise ImportError(
              "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
  
      self.numel = numel
      self.data = torch.zeros(
          self.numel,
          dtype=torch.float,
          device=torch.cuda.current_device(),
          requires_grad=False,
      )
  
      self.chunk_size_mb = chunk_size_mb
      if self.chunk_size_mb > 0:
          chunk_size_bytes = chunk_size_mb * 1024 * 1024
          self.chunk_size_numel = chunk_size_bytes // 4
          self.num_chunks = self.numel // self.chunk_size_numel
          self.numel_per_chunk = [
              self.chunk_size_numel
          ] * self.num_chunks
          if self.numel % self.chunk_size_numel != 0:
              self.num_chunks += 1
              self.numel_per_chunk.append(
                  self.numel % self.chunk_size_numel
              )
  
          self.start_index_per_chunk = torch.cumsum(
              torch.tensor([0] + self.numel_per_chunk[:-1]), dim=0
          )
          self.current_chunk = 0
          self.computed_numel_per_chunk = [0] * self.num_chunks
  
imports: from contextlib import contextmanager
  import torch
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new GradBucket object from inherited base classes: object with example: obj = GradBucket(numel, chunk_size_mb)