code: def forward(
      ctx,
      oldx,
      in_weight,
      in_bias,
      out_weight,
      width,
      depth,
      parallel_size,
      n_nodes,
  ):
      # oldx has shape (..., width)
      x = oldx.reshape(-1, width)
      # x has shape (batch_size, width)
  
      new_logits = fff_cuda.forward(
          x,
          in_weight,
          in_bias,
          out_weight,
          width,
          depth,
          parallel_size,
          n_nodes,
      )
  
      ret = new_logits.reshape_as(oldx)
      return ret
  
file_path: bert-ultrafast/benchmark_cuda/fff_cuda/fff.py
imports: import math
  from torch import nn
  from torch.autograd import Function
  import torch
  import fff_cuda
label_desc: Use image node_id 1 for calling the FFFFunction obj's underlying member method code with example usage: obj.forward(ctx, oldx, in_weight, in_bias, out_weight, width, depth, parallel_size, n_nodes) and returns: ret