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