Update modeling_autoencoder.py

modeling_autoencoder.py

@@ -206,10 +206,13 @@ def create_layers(
         elif layer_types in ['lstm', 'rnn', 'gru']:
             rnn_layer = getattr(nn, layer_types.upper())  # Dynamically get the RNN layer class.
             half_output_dim = output_dim // (2 if bidirectional else 1)
-            if idx == 0:
+            if model_section == "decoder":
+                if idx == 0:
+                    layers.append(rnn_layer(input_dim, half_output_dim, batch_first=True, bidirectional=bidirectional))
+                else: 
+                    layers.append(rnn_layer(input_dim*2, half_output_dim, batch_first=True, bidirectional=bidirectional))
+            else:
                 layers.append(rnn_layer(input_dim, half_output_dim, batch_first=True, bidirectional=bidirectional))
-            else: 
-                layers.append(rnn_layer(input_dim*2, half_output_dim, batch_first=True, bidirectional=bidirectional))
         # Add dropout layer between layers, except for the last layer.
         if (idx != num_layers - 1) and (dropout_rate is not None):
             layers.append(nn.Dropout(dropout_rate))