is there a way to implement the TimeRepresentationAgent into this iteration of my model?
import torch
import torch.nn as nn
import numpy as np

Chomp1d remains unchanged
class Chomp1d(nn.Module):
def init(self, chomp_size):
super(Chomp1d, self).init()
self.chomp_size = chomp_size


def forward(self, x):
    if self.chomp_size > 0:
        return x[:, :, :-self.chomp_size]
    else:
        return x
Residual Block remains unchanged
class ResidualBlock(nn.Module):
def init(self, in_channels, out_channels, kernel_size, dilation, dropout):
super(ResidualBlock, self).init()
padding = (kernel_size - 1) * dilation


    self.conv1 = nn.utils.weight_norm(
        nn.Conv1d(in_channels, out_channels, kernel_size,
                  padding=padding, dilation=dilation)
    )
    self.chomp1 = Chomp1d(padding)
    self.relu1 = nn.ReLU()
    self.dropout1 = nn.Dropout(dropout)
    
    self.conv2 = nn.utils.weight_norm(
        nn.Conv1d(out_channels, out_channels, kernel_size,
                  padding=padding, dilation=dilation)
    )
    self.chomp2 = Chomp1d(padding)
    self.relu2 = nn.ReLU()
    self.dropout2 = nn.Dropout(dropout)

    # Adjust with a 1x1 convolution if in_channels != out_channels
    self.downsample = nn.Conv1d(in_channels, out_channels, kernel_size=1) \
                      if in_channels != out_channels else None

    self.relu = nn.ReLU()
    self.init_weights()
    
def init_weights(self):
    nn.init.normal_(self.conv1.weight, std=0.01)
    nn.init.normal_(self.conv2.weight, std=0.01)
    if self.downsample is not None:
        nn.init.normal_(self.downsample.weight, std=0.01)
        
def forward(self, x):
    residual = x  # Save the input for the residual connection
    out = self.conv1(x)
    out = self.chomp1(out)  # Remove extra padding
    out = self.relu1(out)
    out = self.dropout1(out)
    out = self.conv2(out)
    out = self.chomp2(out)  # Remove extra padding
    out = self.relu2(out)
    out = self.dropout2(out)
    if self.downsample is not None:
        residual = self.downsample(residual)
    return self.relu(out + residual)  # Residual connection with addition
class TCNEncoder(nn.Module):
def init(self, input_size, hidden_size, num_layers, kernel_size, dropout):
super(TCNEncoder, self).init()
layers = []
in_channels = input_size


    for i in range(num_layers):
        dilation = 2 ** i
        out_channels = hidden_size
        layers.append(ResidualBlock(in_channels, out_channels, kernel_size, dilation, dropout))
        in_channels = out_channels  # For next layer

    self.network = nn.Sequential(*layers)

def forward(self, x):
    # x shape: (batch_size, seq_length, input_size)
    x = x.transpose(1, 2)  # Change to (batch_size, input_size, seq_length)
    x = self.network(x)
    x = x.transpose(1, 2)  # Back to (batch_size, seq_length, hidden_size)
    return x
class TransformerDecoder(nn.Module):
def init(self, input_size, hidden_size, num_layers, num_heads, dropout):
super(TransformerDecoder, self).init()
self.embedding = nn.Linear(input_size, hidden_size)
self.positional_encoding = PositionalEncoding(hidden_size, dropout)
decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_size, nhead=num_heads, dropout=dropout)
self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
# self.fc_out = nn.Linear(hidden_size, input_size) # Not needed here


def forward(self, tgt, memory, tgt_mask):
    # tgt: (batch_size, tgt_seq_length, input_size)
    tgt_emb = self.embedding(tgt)
    tgt_emb = self.positional_encoding(tgt_emb)
    memory = memory.transpose(0, 1)  # (src_seq_length, batch_size, hidden_size)
    tgt_emb = tgt_emb.transpose(0, 1)  # (tgt_seq_length, batch_size, hidden_size)
    output = self.transformer_decoder(tgt_emb, memory, tgt_mask=tgt_mask)
    output = output.transpose(0, 1)  # Back to (batch_size, tgt_seq_length, hidden_size)
    return output
#parallel model
class TCNEncoderTransformerDecoderEnsemble(nn.Module):
def init(self, input_size, hidden_size, num_layers, kernel_size,
num_heads, dropout, subset_size):
super(TCNEncoderTransformerDecoderEnsemble, self).init()


    self.input_size = input_size
    self.subset_size = subset_size
    self.hidden_size = hidden_size

    # Calculate the number of feature subsets
    self.num_subsets = (input_size + subset_size - 1) // subset_size  # Ceiling division

    # Create sub-models for each subset
    self.encoders = nn.ModuleList()
    self.decoders = nn.ModuleList()
    self.attention_layers = nn.ModuleList()
    self.final_fc_layers = nn.ModuleList()

    # Shared positional encoding
    self.positional_encoding = PositionalEncoding(hidden_size, dropout)
    self.dropout = nn.Dropout(dropout)

    # Separate alpha parameters for each subset (optional)
    self.alphas = nn.ParameterList([nn.Parameter(torch.full((hidden_size,), 0.5)) for _ in range(self.num_subsets)])

    # Create mini-models for each feature subset
    for i in range(self.num_subsets):
        # Calculate the size of the current subset
        current_input_size = min(subset_size, input_size - i * subset_size)

        # Initialize sub-models
        encoder = TCNEncoder(current_input_size, hidden_size, num_layers, kernel_size, dropout)
        decoder = TransformerDecoder(current_input_size, hidden_size, num_layers, num_heads, dropout)
        attention_layer = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=num_heads, dropout=dropout)
        final_fc_layer = nn.Linear(hidden_size, current_input_size)

        self.encoders.append(encoder)
        self.decoders.append(decoder)
        self.attention_layers.append(attention_layer)
        self.final_fc_layers.append(final_fc_layer)

def forward(self, src, tgt):
    # src: (batch_size, src_seq_length, input_size)
    # tgt: (batch_size, tgt_seq_length, input_size)

    batch_size = src.size(0)
    src_seq_len = src.size(1)
    tgt_seq_len = tgt.size(1)

    # Split the source and target into subsets
    src_splits = torch.chunk(src, self.num_subsets, dim=2)
    tgt_splits = torch.chunk(tgt, self.num_subsets, dim=2)

    # Initialize lists to store outputs
    encoder_outputs = []
    decoder_outputs = []
    attn_outputs = []
    final_outputs = []

    # Process all subsets in parallel using list comprehensions
    for i in range(self.num_subsets):
        # Extract feature subsets
        src_sub = src_splits[i]
        tgt_sub = tgt_splits[i]

        # Get the corresponding sub-models
        encoder = self.encoders[i]
        decoder = self.decoders[i]
        attention_layer = self.attention_layers[i]
        final_fc_layer = self.final_fc_layers[i]
        alpha = torch.sigmoid(self.alphas[i]).to(src.device).view(1, 1, -1)

        # Encoder
        memory = encoder(src_sub)  # (batch_size, src_seq_len, hidden_size)
        memory = self.positional_encoding(memory)  # Add positional encoding

        # Decoder
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_sub.size(1)).to(tgt_sub.device)
        decoder_output = decoder(tgt_sub, memory, tgt_mask)  # (batch_size, tgt_seq_len, hidden_size)

        # Reshape for MultiheadAttention
        decoder_output_t = decoder_output.transpose(0, 1)  # (tgt_seq_len, batch_size, hidden_size)
        memory_t = memory.transpose(0, 1)  # (src_seq_len, batch_size, hidden_size)

        # Additional Attention Layer
        attn_output, _ = attention_layer(decoder_output_t, memory_t, memory_t)
        attn_output = attn_output.transpose(0, 1)  # (batch_size, tgt_seq_len, hidden_size)

        # Fusion Strategy 2: Weighted Sum
        combined = alpha * decoder_output + (1 - alpha) * attn_output
        combined = self.dropout(combined)

        # Final Output Layer
        final_output_sub = final_fc_layer(combined)  # (batch_size, tgt_seq_len, subset_size)
        final_outputs.append(final_output_sub)

    # Concatenate outputs along the feature dimension
    final_output = torch.cat(final_outputs, dim=2)  # (batch_size, tgt_seq_len, input_size)

    return final_output
Function to train the model with teacher forcing
def train_model(model, train_loader, val_loader, num_epochs, learning_rate, gradient_clip_val, train_only = False):
criterion = nn.MSELoss()
optimizer = torch.optim.AdamW([
{'params': model.parameters(), 'weight_decay' : 0.01},
#{'params': criterion.parameters(), 'weight_decay' : 0.0}
], lr=learning_rate)
processor_type= 'cuda' if torch.cuda.is_available() else 'cpu'
scaler = GradScaler()
device = torch.device(processor_type)
model.to(device)


#scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=learning_rate*2, steps_per_epoch=len(train_loader), epochs=num_epochs)

# Track best validation loss for early stopping
best_val_loss = float('inf')
patience = 10
patience_counter = 0
pbar = tqdm(range(num_epochs), desc=f"Epochs", unit="epoch")
for epoch in pbar:
    model.train()
    train_loss = 0.0
    train_sMAPE = 0.0
    train_MASE = 0.0
    for batch_x, batch_y in train_loader:
        
        batch_x, batch_y = _prepare_batch(batch_x, batch_y, device)
        # Prepare decoder input by shifting targets (teacher forcing)
        decoder_input = _prepare_decoder_input(batch_x, batch_y, True)

        optimizer.zero_grad(set_to_none=True)
        
        with autocast(device_type= processor_type):
            output = model(batch_x, decoder_input)
            loss = criterion(output, batch_y)
            batch_sMAPE = smape(batch_y, output)
            batch_MASE = 0#mase(batch_y, output, torch.tensor(train_data_Weekly.values, dtype=torch.float32).to(device))
        # Scale the loss and call backward
        
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clip_val)
        scaler.step(optimizer)
        scaler.update()
        train_loss += loss.item()
        train_sMAPE += batch_sMAPE
        train_MASE += batch_MASE
        #scheduler.step()
        
    train_loss /= len(train_loader)
    train_sMAPE /= len(train_loader)
    train_MASE /= len(train_loader)

    if not(train_only):
        Flag, best_val_loss, losses = validate_model(model, val_loader, criterion, best_val_loss, processor_type, patience_counter, patience)
        val_loss, val_sMAPE, val_MASE = losses

        #break if patience runs out
        if not(Flag):
            break
        
    # Log to wandb
    wandb.log({"Epoch": epoch + 1, "Train Loss": train_loss, "Train sMAPE": train_sMAPE, "Train MASE": train_MASE, 
               "Validation Loss": val_loss, "Validation sMAPE": val_sMAPE, "Validation MASE": val_MASE,
               #"Learning Rate": scheduler.get_last_lr()[0]
               #"Loss alpha": criterion.logits[0],"Loss beta": criterion.logits[1],"Loss gamma": criterion.logits[2]
               #,"Loss dealta": criterion.logits[3],"Loss lambda": criterion.logits[4]
              })
    # Update progress bar
    pbar.set_description(f"Epoch {epoch + 1}/{num_epochs}" 
    +f"- Average Loss: {val_loss:.4f}" 
    #+f"- LR: {scheduler.get_last_lr()[0]:.6f}"
    )

# Finalize wandb logging
#wandb.finish()
return model
def validate_model(model, val_loader, criterion, best_val_loss, processor_type, patience_counter, patience):
model.eval()
val_loss = 0.0
val_sMAPE = 0.0
val_MASE = 0.0
with torch.no_grad():
for batch_x, batch_y in val_loader:
if batch_y.dim() == 1:
batch_x = batch_x.to(device).unsqueeze(-1)
batch_y = batch_y.to(device).unsqueeze(-1)
elif batch_y.dim() > 1:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)
else:
print("The input doesn't fit the expected shape of a uni or multi variate Timeseries")
decoder_input = torch.zeros_like(batch_y)
decoder_input[:, 1:, :] = batch_y[:, :-1, :]
decoder_input[:, 0, :] = batch_x[:, -1, :]
with autocast(device_type= processor_type):
output = model(batch_x, decoder_input)
loss = criterion(output, batch_y)
batch_sMAPE = smape(batch_y, output)
val_loss += loss.item()
val_sMAPE += batch_sMAPE
#val_MASE += batch_MASE


    val_loss /= len(val_loader)
    val_sMAPE /= len(val_loader)
    val_MASE /= len(val_loader)
        # Early stopping check
if val_loss < best_val_loss:
    best_val_loss = val_loss
    patience_counter = 0
elif val_loss > (best_val_loss*1.025):
    patience_counter += 1
    if patience_counter >= patience:
        print(f"\nEarly stopping triggered after {epoch + 1} epochs")
        return False
return True, best_val_loss, (val_loss, val_sMAPE, val_MASE)