code: def __init__(
      self,
      ctx: Context,
      model: HeteroNNModelHost,
      training_args: TrainingArguments,
      train_set: Dataset,
      val_set: Dataset = None,
      optimizer=None,
      data_collator: Callable = None,
      scheduler=None,
      tokenizer: Optional[PreTrainedTokenizer] = None,
      callbacks: Optional[List[TrainerCallback]] = [],
      compute_metrics: Optional[
          Callable[[EvalPrediction], Dict]
      ] = None,
  ):
      assert isinstance(model, HeteroNNModelHost), (
          "Model should be a HeteroNNModelHost instance, " "but got {}."
      ).format(type(model))
  
      if model.need_mpc_init():
          ctx.mpc.init()
  
      model.setup(ctx=ctx)
      super().__init__(
          ctx=ctx,
          model=model,
          training_args=training_args,
          train_set=train_set,
          val_set=val_set,
          loss_fn=None,
          optimizer=optimizer,
          data_collator=data_collator,
          scheduler=scheduler,
          tokenizer=tokenizer,
          callbacks=callbacks,
          compute_metrics=compute_metrics,
      )
  
imports: import torch
  from dataclasses import dataclass, fields
  from enum import Enum
  from torch import nn
  from typing import Any, Dict, List, Union, Callable, Literal
  from fate.arch import Context
  from torch.utils.data import Dataset
  from transformers import PreTrainedTokenizer
  from transformers import EvalPrediction
  from transformers.trainer_callback import TrainerCallback
  from typing import Optional
  from fate.ml.nn.model_zoo.hetero_nn_model import HeteroNNModelGuest, HeteroNNModelHost
  from fate.ml.nn.trainer.trainer_base import HeteroTrainerBase, TrainingArguments
  from fate.ml.nn.model_zoo.hetero_nn_model import StdAggLayerArgument, FedPassArgument, SSHEArgument
  from fate.ml.nn.model_zoo.hetero_nn_model import TopModelStrategyArguments
label_desc: Use image node_id 1 to create a new HeteroNNTrainerHost object from inherited base classes: HeteroTrainerBase with example: obj = HeteroNNTrainerHost(ctx, model, training_args, train_set, val_set, optimizer, data_collator, scheduler, tokenizer, callbacks, compute_metrics)