code: def compute_memory(module, inp, out):
      if isinstance(module, (nn.ReLU, nn.ReLU6, nn.ELU, nn.LeakyReLU)):
          return compute_ReLU_memory(module, inp[0], out[0])
      elif isinstance(module, nn.PReLU):
          return compute_PReLU_memory(module, inp[0], out[0])
      elif isinstance(module, nn.Conv2d):
          return compute_Conv2d_memory(module, inp[0], out[0])
      elif isinstance(module, nn.BatchNorm2d):
          return compute_BatchNorm2d_memory(module, inp[0], out[0])
      elif isinstance(module, nn.Linear):
          return compute_Linear_memory(module, inp[0], out[0])
      elif isinstance(module, (nn.AvgPool2d, nn.MaxPool2d)):
          return compute_Pool2d_memory(module, inp[0], out[0])
      else:
          # print(f"[Memory]: {type(module).__name__} is not supported!")
          return 0, 0
      pass
  
file_path: tensorwatch/tensorwatch/model_graph/torchstat/compute_memory.py
imports: import torch.nn
  import torch
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: compute_memory(module, inp, out) and returns: compute_ReLU_memory, compute_PReLU_memory, compute_Conv2d_memory, compute_BatchNorm2d_memory, compute_Linear_memory, compute_Pool2d_memory, int,  int