code: def _profile_concrete(
      target: Callable, *args, **kwargs
  ) -> Tuple[Tuple[Any, ...], GraphInfo]:
      """Profile a Callable function with args and kwargs on concrete devices by https://github.com/Cypher30
      To profile the actual forward memory, we first run target in the context torch.no_grad() to get
      the fwd_mem_out, then we run target with grad enable to found the extra memory stored in the memory
      by memory allocated minus the fwd_mem_out.
      To profile the actual backward memory, we first make dummy gradient for torch.autograd.backward, then
      find the bwd_mem_tmp with memory peak during the process minus bwd_mem_out(it is actually equal to size
      of args and kwargs).
      We also add time stamps to profile the real forward and backward time.
  
      Args:
          target (Callable): A Callable function
          args (Any): Arguments
          kwargs (Any): Arguments
  
      Returns:
          Tuple[Tuple[Any, ...], GraphInfo]: Output for next node & memory cost and real forward and backward
          time.
      """
  
      graphinfo = GraphInfo()
  
      # detach input from the graph
      args = tree_map(detach_variables, args)
      kwargs = tree_map(detach_variables, kwargs)
      if isinstance(target, str):
          # args[0] is the `self` object for this method call
          self_obj, *args_tail = args
  
          # calculate fwd_mem_out
          mem_stamp0 = torch.cuda.memory_allocated()
          with torch.no_grad():
              out = getattr(self_obj, target)(*args_tail, **kwargs)
          mem_stamp1 = torch.cuda.memory_allocated()
          graphinfo.fwd_mem_out = mem_stamp1 - mem_stamp0
          del out
  
          # calculate fwd_mem_tmp & fwd_time
          mem_stamp0 = torch.cuda.memory_allocated()
          fwd_time0 = time.time()
          out = getattr(self_obj, target)(*args_tail, **kwargs)
          fwd_time1 = time.time()
          graphinfo.fwd_time = fwd_time1 - fwd_time0
          mem_stamp1 = torch.cuda.memory_allocated()
          graphinfo.fwd_mem_tmp = (
              mem_stamp1 - mem_stamp0 - graphinfo.fwd_mem_out
          )
  
          # calculate bwd_mem_tmp & bwd_time
          grad_tensors = tree_map(
              lambda x: torch.ones_like(x)
              if isinstance(x, torch.Tensor)
              else None,
              out,
          )
          torch.cuda.reset_peak_memory_stats()
          mem_stamp0 = torch.cuda.memory_allocated()
          bwd_time0 = time.time()
          torch.autograd.backward(out, grad_tensors=grad_tensors)
          bwd_time1 = time.time()
          graphinfo.bwd_time = bwd_time1 - bwd_time0
          mem_stamp1 = torch.cuda.max_memory_allocated()
  
          # calculate bwd memory stats
          # NOTE: the module should add param to bwd_mem_out for bwd_mem_tmp calculation
          graphinfo.bwd_mem_out = activation_size(
              args
          ) + activation_size(kwargs)
          graphinfo.bwd_mem_out += (
              parameter_size(target.__self__)
              if hasattr(target.__self__, "parameters")
              else 0
          )
          graphinfo.bwd_mem_tmp = (
              mem_stamp1 - mem_stamp0 - graphinfo.bwd_mem_out
          )
  
      else:
          # calculate fwd_mem_out
          mem_stamp0 = torch.cuda.memory_allocated()
          with torch.no_grad():
              out = target(*args, **kwargs)
          mem_stamp1 = torch.cuda.memory_allocated()
          graphinfo.fwd_mem_out = mem_stamp1 - mem_stamp0
          del out
  
          # calculate fwd_mem_tmp & fwd_time
          mem_stamp0 = torch.cuda.memory_allocated()
          fwd_time0 = time.time()
          out = target(*args, **kwargs)
          fwd_time1 = time.time()
          graphinfo.fwd_time = fwd_time1 - fwd_time0
          mem_stamp1 = torch.cuda.memory_allocated()
          graphinfo.fwd_mem_tmp = (
              mem_stamp1 - mem_stamp0 - graphinfo.fwd_mem_out
          )
  
          # calculate bwd_mem_tmp & bwd_time
          grad_tensors = tree_map(
              lambda x: torch.ones_like(x)
              if isinstance(x, torch.Tensor)
              else None,
              out,
          )
          torch.cuda.reset_peak_memory_stats()
          mem_stamp0 = torch.cuda.memory_allocated()
          bwd_time0 = time.time()
          torch.autograd.backward(out, grad_tensors=grad_tensors)
          bwd_time1 = time.time()
          graphinfo.bwd_time = bwd_time1 - bwd_time0
          mem_stamp1 = torch.cuda.max_memory_allocated()
  
          # calculate bwd memory stats
          # NOTE: the module should add param to bwd_mem_out for bwd_mem_tmp calculation
          graphinfo.bwd_mem_out = activation_size(
              args
          ) + activation_size(kwargs)
          graphinfo.bwd_mem_out += (
              parameter_size(target.__self__)
              if hasattr(target.__self__, "parameters")
              else 0
          )
          graphinfo.bwd_mem_tmp = (
              mem_stamp1 - mem_stamp0 - graphinfo.bwd_mem_out
          )
  
      return tree_map(detach_variables, out), graphinfo
  
imports: import time
  from functools import partial
  from typing import Any, Callable, Dict, Tuple
  import torch
  from torch.fx import Graph, Node
  from torch.fx.node import Argument, Target
  from torch.nn.parameter import Parameter
  from torch.utils._pytree import tree_map
  from .._compatibility import compatibility
  from .constants import ALIAS_ATEN, OUTPUT_SAVED_MOD, OUTPUT_SAVED_OPS
  from .dataflow import GraphInfo, Phase, autograd_graph_analysis, is_phase
  from .memory_utils import activation_size, parameter_size
  from .opcount import flop_mapping
  from .tensor import MetaTensor
label_desc: Use image node_id 4 for calling a global function with example usage: _profile_concrete(target) and returns: tree_map,  graphinfo