code: def softmax(x, dim: int, onnx_trace: bool = False):
      if onnx_trace:
          return F.softmax(x.float(), dim=dim)
      else:
          return F.softmax(x, dim=dim, dtype=torch.float32)
  
file_path: lightseq/lightseq/training/ops/pytorch/util.py
imports: import math
  from typing import Callable
  import numpy
  import torch
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: softmax(x, dim, onnx_trace) and returns: F, F