code: def __init__(
      self,
      model_config,
      tra_config,
      model_type="LSTM",
      lr=1e-3,
      n_epochs=500,
      early_stop=50,
      smooth_steps=5,
      max_steps_per_epoch=None,
      freeze_model=False,
      model_init_state=None,
      lamb=0.0,
      rho=0.99,
      seed=None,
      logdir=None,
      eval_train=True,
      eval_test=False,
      avg_params=True,
      **kwargs,
  ):
      np.random.seed(seed)
      torch.manual_seed(seed)
  
      self.logger = get_module_logger("TRA")
      self.logger.info("TRA Model...")
  
      self.model = eval(model_type)(**model_config).to(device)
      if model_init_state:
          self.model.load_state_dict(
              torch.load(model_init_state, map_location="cpu")["model"]
          )
      if freeze_model:
          for param in self.model.parameters():
              param.requires_grad_(False)
      else:
          self.logger.info(
              "# model params: %d"
              % sum([p.numel() for p in self.model.parameters()])
          )
  
      self.tra = TRA(self.model.output_size, **tra_config).to(device)
      self.logger.info(
          "# tra params: %d"
          % sum([p.numel() for p in self.tra.parameters()])
      )
  
      self.optimizer = optim.Adam(
          list(self.model.parameters()) + list(self.tra.parameters()),
          lr=lr,
      )
  
      self.model_config = model_config
      self.tra_config = tra_config
      self.lr = lr
      self.n_epochs = n_epochs
      self.early_stop = early_stop
      self.smooth_steps = smooth_steps
      self.max_steps_per_epoch = max_steps_per_epoch
      self.lamb = lamb
      self.rho = rho
      self.seed = seed
      self.logdir = logdir
      self.eval_train = eval_train
      self.eval_test = eval_test
      self.avg_params = avg_params
  
      if self.tra.num_states > 1 and not self.eval_train:
          self.logger.warn(
              "`eval_train` will be ignored when using TRA"
          )
  
      if self.logdir is not None:
          if os.path.exists(self.logdir):
              self.logger.warn(f"logdir {self.logdir} is not empty")
          os.makedirs(self.logdir, exist_ok=True)
  
      self.fitted = False
      self.global_step = -1
  
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 1 to create a new TRAModel object from inherited base classes: Model with example: obj = TRAModel(model_config, tra_config, model_type, lr, n_epochs, early_stop, smooth_steps, max_steps_per_epoch, freeze_model, model_init_state, lamb, rho, seed, logdir, eval_train, eval_test, avg_params)