How do I implement the following MultiScaleFeatureFusion into my current model:
class MultiScaleFeatureFusion(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super(MultiScaleFeatureFusion, self).__init__()
        self.attention_layers = nn.ModuleList([
            nn.MultiheadAttention(hidden_size, num_heads) for _ in range(3)
        ])
        self.norm_layers = nn.ModuleList([
            nn.LayerNorm(hidden_size) for _ in range(3)
        ])
        self.fusion_gate = nn.Linear(hidden_size * 3, 3)

    def forward(self, decoder_output, memory):
        # Multi-scale attention
        outputs = []
        for attn, norm in zip(self.attention_layers, self.norm_layers):
            normed_decoder = norm(decoder_output.transpose(0, 1))
            attn_output, _ = attn(normed_decoder, memory, memory)
            outputs.append(attn_output.transpose(0, 1))
        
        # Adaptive fusion with gating mechanism
        concat_features = torch.cat(outputs, dim=-1)
        gates = self.fusion_gate(concat_features).softmax(dim=-1)
        
        # Weighted combination
        final_output = sum(g * o for g, o in zip(gates.unbind(-1), outputs))
        return final_output


My current Model:
# Weight normzlizatioin 22~23, Layer 24~25.

# Chomp1d removes extra time steps from the end of the sequence
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 with causal convolutions and Chomp1d
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)
        # If the input and output channels are not the same, adjust with a 1x1 convolution
        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

# Updated TCN Encoder with Residual Blocks 0.4
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

    
# Transformer Decoder with masked multi-head attention
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)

    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

    
# TCN-Encoder-Transformer-Decoder Model
class TCNEncoderTransformerDecoder(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, kernel_size, num_heads, dropout):
        super(TCNEncoderTransformerDecoder, self).__init__()
        self.encoder = TCNEncoder(input_size, hidden_size, num_layers, kernel_size, dropout)
        # Additional Attention Layer (Fusion Cross-Attention) ver = 0.4
        self.attention_layer = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=num_heads, dropout=dropout)
        self.decoder = TransformerDecoder(input_size, hidden_size, num_layers, num_heads, dropout)
        self.positional_encoding = PositionalEncoding(hidden_size, dropout) # Encoding Strategy 1
        #self.positional_encoding = LearnablePositionalEncoding(hidden_size)  # For encoder outputs

        # Fusion Layer ver = 0.4 if Fusion startgy 1 then multiply hidden_size by 2
        self.final_fc_layer = nn.Linear(hidden_size, input_size)
        self.dropout = nn.Dropout(dropout)

        # Alpha parameter for weighted sum fusion ver = 0.4
        # Initialize alpha as a learnable parameter with value 0.5
        self.alpha = nn.Parameter(torch.full((hidden_size,), 0.5))

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

        # Encoder
        memory = self.encoder(src)  # (batch_size, src_seq_length, hidden_size)
        memory = self.positional_encoding(memory)  # Add positional encoding to memory

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

        # Reshape for MultiheadAttention (seq_length, batch_size, embed_dim)
        decoder_output_t = decoder_output.transpose(0, 1)  # (tgt_seq_length, batch_size, hidden_size)
        memory_t = memory.transpose(0, 1)  # (src_seq_length, batch_size, hidden_size)

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

        # Fusion Strategy 1: Concatenate decoder output and attention output
        #combined = torch.cat((decoder_output, attn_output), dim=-1)  # (batch_size, tgt_seq_length, hidden_size * 2)
        #combined = self.dropout(combined)#

        # Fusion Strategy 2: Weighted Sum of decoder output and attention output
        # I must always ensure alpha is in [0,1] using sigmoid or similar functions
        # Tested scaler alpha and it perfromed worse than 4% worse than Strategy 1 so switched it to vector
        #alpha = torch.sigmoid(self.alpha).to(decoder_output.device).view(1, 1, -1)  # Shape: (1, 1, hidden_size)
        #combined = alpha * decoder_output + (1 - alpha) * attn_output  # (batch_size, tgt_seq_length, hidden_size)
        #combined = self.dropout(combined)

        # Fusion Strategy 2: Multiplication
        combined = decoder_output * attn_output
        combined = self.dropout(combined)
        
        # Final Output Layer
        final_output = self.final_fc_layer(combined)  # (batch_size, tgt_seq_length, input_size)

        return final_output


# Positional Encoding class (Encoding Strategy 1)
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)  # Create positional encoding matrix
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)  # Apply sine to even indices
        pe[:, 1::2] = torch.cos(position * div_term)  # Apply cosine to odd indices
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)  # Save positional encoding matrix

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :].to(x.device)  # Add positional encoding
        return self.dropout(x)


# Learnable Positional Encoding class (Encoding Stratgy 2)
class LearnablePositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(LearnablePositionalEncoding, self).__init__()
        self.pe = nn.Parameter(torch.randn(1, max_len, d_model))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :].to(x.device)
        return x