code: def torch_bmm(a, b):
      if a.device.type == "cpu":
          return torch_bmm_old_(a, b)
  
      batch, A_rows, common = a.size()
      B_cols = b.size(2)
      C = torch.empty((batch, A_rows, B_cols), device=a.device)
      if a.size(1) == 1:
          torch_bmm_old_(a, b, out=C)
      else:
          linear_q4_0.bmm(a.contiguous(), b.contiguous(), C)
      return C
  
imports: import torch
  import linear_q4_0
label_desc: Use image node_id 1 for calling a global function with example usage: torch_bmm(a, b) and returns: C, torch_bmm_old_