Update modeling_autoencoder.py

modeling_autoencoder.py

@@ -304,10 +304,11 @@ class AutoEncoder(PreTrainedModel):
             for layer in self.encoder:
                 if isinstance(layer, nn.LSTM):
                     input_ids, (h_n, c_n) = layer(input_ids)
-                    input_ids = input_ids.flatten_parameters()
+                    print(input_ids.size())
+
                 elif isinstance(layer, nn.RNN) or isinstance(layer, nn.GRU):
                     input_ids, h_o = layer(input_ids)
-                    input_ids = input_ids.flatten_parameters()
+
                 else:
                     input_ids = layer(input_ids)
             # Hidden Vector
@@ -316,10 +317,11 @@ class AutoEncoder(PreTrainedModel):
             for layer in self.decoder:
                 if isinstance(layer, nn.LSTM):
                     input_ids, (h_n, c_n) = layer(input_ids)
-                    input_ids = input_ids.flatten_parameters()
+                    print(input_ids.size())
+
                 elif isinstance(layer, nn.RNN) or isinstance(layer, nn.GRU):
                     input_ids, h_o = layer(input_ids)
-                    input_ids = input_ids.flatten_parameters()
+
                 else:
                     input_ids = layer(input_ids)