code: def fit(self, dataset, evals_result=dict()):
      train_set, valid_set, test_set = dataset.prepare(
          ["train", "valid", "test"]
      )
  
      best_score = -1
      best_epoch = 0
      stop_rounds = 0
      best_params = {
          "model": copy.deepcopy(self.model.state_dict()),
          "tra": copy.deepcopy(self.tra.state_dict()),
      }
      params_list = {
          "model": collections.deque(maxlen=self.smooth_steps),
          "tra": collections.deque(maxlen=self.smooth_steps),
      }
      evals_result["train"] = []
      evals_result["valid"] = []
      evals_result["test"] = []
  
      # train
      self.fitted = True
      self.global_step = -1
  
      if self.tra.num_states > 1:
          self.logger.info("init memory...")
          self.test_epoch(train_set)
  
      for epoch in range(self.n_epochs):
          self.logger.info("Epoch %d:", epoch)
  
          self.logger.info("training...")
          self.train_epoch(train_set)
  
          self.logger.info("evaluating...")
          # average params for inference
          params_list["model"].append(
              copy.deepcopy(self.model.state_dict())
          )
          params_list["tra"].append(
              copy.deepcopy(self.tra.state_dict())
          )
          self.model.load_state_dict(
              average_params(params_list["model"])
          )
          self.tra.load_state_dict(average_params(params_list["tra"]))
  
          # NOTE: during evaluating, the whole memory will be refreshed
          if self.tra.num_states > 1 or self.eval_train:
              train_set.clear_memory()  # NOTE: clear the shared memory
              train_metrics = self.test_epoch(train_set)[0]
              evals_result["train"].append(train_metrics)
              self.logger.info("\ttrain metrics: %s" % train_metrics)
  
          valid_metrics = self.test_epoch(valid_set)[0]
          evals_result["valid"].append(valid_metrics)
          self.logger.info("\tvalid metrics: %s" % valid_metrics)
  
          if self.eval_test:
              test_metrics = self.test_epoch(test_set)[0]
              evals_result["test"].append(test_metrics)
              self.logger.info("\ttest metrics: %s" % test_metrics)
  
          if valid_metrics["IC"] > best_score:
              best_score = valid_metrics["IC"]
              stop_rounds = 0
              best_epoch = epoch
              best_params = {
                  "model": copy.deepcopy(self.model.state_dict()),
                  "tra": copy.deepcopy(self.tra.state_dict()),
              }
          else:
              stop_rounds += 1
              if stop_rounds >= self.early_stop:
                  self.logger.info("early stop @ %s" % epoch)
                  break
  
          # restore parameters
          self.model.load_state_dict(params_list["model"][-1])
          self.tra.load_state_dict(params_list["tra"][-1])
  
      self.logger.info(
          "best score: %.6lf @ %d" % (best_score, best_epoch)
      )
      self.model.load_state_dict(best_params["model"])
      self.tra.load_state_dict(best_params["tra"])
  
      metrics, preds = self.test_epoch(test_set, return_pred=True)
      self.logger.info("test metrics: %s" % metrics)
  
      if self.logdir:
          self.logger.info("save model & pred to local directory")
  
          pd.concat(
              {
                  name: pd.DataFrame(evals_result[name])
                  for name in evals_result
              },
              axis=1,
          ).to_csv(self.logdir + "/logs.csv", index=False)
  
          torch.save(best_params, self.logdir + "/model.bin")
  
          preds.to_pickle(self.logdir + "/pred.pkl")
  
          info = {
              "config": {
                  "model_config": self.model_config,
                  "tra_config": self.tra_config,
                  "lr": self.lr,
                  "n_epochs": self.n_epochs,
                  "early_stop": self.early_stop,
                  "smooth_steps": self.smooth_steps,
                  "max_steps_per_epoch": self.max_steps_per_epoch,
                  "lamb": self.lamb,
                  "rho": self.rho,
                  "seed": self.seed,
                  "logdir": self.logdir,
              },
              "best_eval_metric": -best_score,  # NOTE: minux -1 for minimize
              "metric": metrics,
          }
          with open(self.logdir + "/info.json", "w") as f:
              json.dump(info, f)
  
imports: import os
  import copy
  import math
  import json
  import collections
  import numpy
  import pandas
  import torch
  import torch.nn
  import torch.optim
  import torch.nn.functional
  from tqdm import tqdm
  from qlib.utils import get_or_create_path
  from qlib.log import get_module_logger
  from qlib.model.base import Model
label_desc: Use image node_id 4 for calling the TRAModel obj's underlying member method code with example usage: obj.fit(dataset, evals_result) without return types