class_docstr: Estimate memory with chunk
class_docstr_tok: ["Estimate","memory","with","chunk"]
docstr: return node size in MB
docstr_tok: ["return","node","size","in","MB"]
code: def _get_node_size(self, x: Node) -> float:
      """
      return node size in MB
      """
      x = x.meta["tensor_meta"]
      if not hasattr(x, "numel"):
          out = sum(
              [
                  i.numel
                  * torch.tensor([], dtype=i.dtype).element_size()
                  for i in x
              ]
          )
      else:
          out = x.numel * torch.tensor([], dtype=x.dtype).element_size()
      out = float(out) / 1024**2
      return out
  
file_path: ColossalAI/colossalai/autochunk/estimate_memory.py
imports: from typing import Dict, List
  import torch
  from torch.fx.node import Node
  from .utils import NodeMgr, get_node_shape, is_non_memory_node
label_desc: Use image node_id 2 for calling the EstimateMemory obj's underlying member method code with example usage: obj._get_node_size(x) and returns: out