code: def forward(ctx, input, weight, bias=None):
      # print("In ZeRO Linear Function")
  
      weight_id = id(weight)
      bias_id = id(bias)
  
      # ctx.save_for_backward(input, weight, bias)
      ctx.save_for_backward(
          input, torch.tensor(weight_id), torch.tensor(bias_id)
      )
  
      tensor_map[weight_id] = weight
      tensor_map[bias_id] = bias
  
      if input.dim() == 2 and bias is not None:
          # fused op is marginally faster
          ret = torch.addmm(bias, input, weight.t())
      else:
          output = input.matmul(weight.t())
          if bias is not None:
              output += bias
          ret = output
  
      return ret
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/runtime/zero/linear.py
imports: import math
  import torch
  from torch import Tensor
  from torch.nn.parameter import Parameter
  from torch.nn import init
  from torch.nn.modules.module import Module
  from deepspeed.runtime.utils import noop_decorator
  from deepspeed import comm
label_desc: Use image node_id 1 for calling the LinearFunctionForZeroStage3 obj's underlying member method code with example usage: obj.forward(ctx, input, weight, bias) and returns: ret