code: def __init__(
      self,
      input_size: int,
      num_layers: int,
      num_cells: int,
      cell_type: str,
      history_length: int,
      context_length: int,
      prediction_length: int,
      dropout_rate: float,
      lags_seq: List[int],
      target_dim: int,
      conditioning_length: int,
      diff_steps: int,
      loss_type: str,
      beta_end: float,
      beta_schedule: str,
      residual_layers: int,
      residual_channels: int,
      dilation_cycle_length: int,
      cardinality: List[int] = [1],
      embedding_dimension: int = 1,
      scaling: bool = True,
      **kwargs,
  ) -> None:
      super().__init__(**kwargs)
      self.target_dim = target_dim
      self.prediction_length = prediction_length
      self.context_length = context_length
      self.history_length = history_length
      self.scaling = scaling
  
      assert len(set(lags_seq)) == len(
          lags_seq
      ), "no duplicated lags allowed!"
      lags_seq.sort()
      self.lags_seq = lags_seq
  
      self.cell_type = cell_type
      rnn_cls = {"LSTM": nn.LSTM, "GRU": nn.GRU}[cell_type]
      self.rnn = rnn_cls(
          input_size=input_size,
          hidden_size=num_cells,
          num_layers=num_layers,
          dropout=dropout_rate,
          batch_first=True,
      )
  
      self.denoise_fn = EpsilonTheta(
          target_dim=target_dim,
          cond_length=conditioning_length,
          residual_layers=residual_layers,
          residual_channels=residual_channels,
          dilation_cycle_length=dilation_cycle_length,
      )
  
      self.diffusion = GaussianDiffusion(
          self.denoise_fn,
          input_size=target_dim,
          diff_steps=diff_steps,
          loss_type=loss_type,
          beta_end=beta_end,
          beta_schedule=beta_schedule,
      )
  
      self.distr_output = DiffusionOutput(
          self.diffusion,
          input_size=target_dim,
          cond_size=conditioning_length,
      )
  
      self.proj_dist_args = self.distr_output.get_args_proj(num_cells)
  
      self.embed_dim = 1
      self.embed = nn.Embedding(
          num_embeddings=self.target_dim, embedding_dim=self.embed_dim
      )
  
      if self.scaling:
          self.scaler = MeanScaler(keepdim=True)
      else:
          self.scaler = NOPScaler(keepdim=True)
  
imports: from torch.nn.modules import loss
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  from gluonts.core.component import validated
  from pts.model import weighted_average
  from pts.modules import GaussianDiffusion, DiffusionOutput, MeanScaler, NOPScaler
  from .epsilon_theta import EpsilonTheta
label_desc: Use image node_id 1 to create a new TimeGradTrainingNetwork object from inherited base classes: nn with example: obj = TimeGradTrainingNetwork(input_size, num_layers, num_cells, cell_type, history_length, context_length, prediction_length, dropout_rate, lags_seq, target_dim, conditioning_length, diff_steps, loss_type, beta_end, beta_schedule, residual_layers, residual_channels, dilation_cycle_length, cardinality, embedding_dimension, scaling)