code: def __init__(
      self,
      model,
      cfg,
      combine_margin,
      cross_entropy,
      data_loader,
      optimizer,
      lr_scheduler=None,
  ):
      super().__init__()
  
      if cfg.fp16:
          self.config.enable_amp(True)
          self.set_grad_scaler(make_grad_scaler())
      elif cfg.scale_grad:
          self.set_grad_scaler(make_static_grad_scaler())
  
      self.config.allow_fuse_add_to_output(True)
      self.config.allow_fuse_model_update_ops(True)
  
      self.model = model
  
      self.cross_entropy = cross_entropy
      self.combine_margin = combine_margin
      self.data_loader = data_loader
      self.add_optimizer(optimizer, lr_sch=lr_scheduler)
  
file_path: insightface/recognition/arcface_oneflow/graph.py
imports: import oneflow
  import oneflow.nn
label_desc: Use image node_id 1 to create a new TrainGraph object from inherited base classes: flow.nn with example: obj = TrainGraph(model, cfg, combine_margin, cross_entropy, data_loader, optimizer, lr_scheduler)