Fix swap_tensors path in _apply for modules that inherit from RNNBase…

… (RNN, GRU, LSTM) (pytorch#122800) Pull Request resolved: pytorch#122800 Approved by: https://github.com/albanD (cherry picked from commit cc12668)
mikaylagawarecki · Apr 1, 2024 · d2648f9 · d2648f9
1 parent 8602990
commit d2648f9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 15 deletions.
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
@@ -213,6 +213,7 @@ def flatten_parameters(self) -> None:
                         self.batch_first, bool(self.bidirectional))
 
     def _apply(self, fn, recurse=True):
+        self._flat_weight_refs = []
         ret = super()._apply(fn, recurse)
 
         # Resets _flat_weights

diff --git a/torch/testing/_internal/common_modules.py b/torch/testing/_internal/common_modules.py
@@ -4288,34 +4288,20 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad
                train_and_eval_differ=True,
                module_inputs_func=partial(module_inputs_torch_nn_RNN_GRU, is_rnn=True),
                module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU,
-               skips=(
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to', active_if=lambda p: p['swap']),
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to_empty', active_if=lambda p: p['swap']),),
                decorators=rnn_gru_lstm_module_info_decorators
                ),
     ModuleInfo(torch.nn.GRU,
                train_and_eval_differ=True,
                module_inputs_func=partial(module_inputs_torch_nn_RNN_GRU, is_rnn=False),
                module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU,
-               skips=(
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to', active_if=lambda p: p['swap']),
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to_empty', active_if=lambda p: p['swap']),),
                decorators=rnn_gru_lstm_module_info_decorators),
     ModuleInfo(torch.nn.LSTM,
                train_and_eval_differ=True,
                module_inputs_func=module_inputs_torch_nn_LSTM,
                module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU,
                skips=(
                    # LSTM with projections is not currently supported with MPS
-                   DecorateInfo(skipMPS),
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to', active_if=lambda p: p['swap']),
-                   # RNNBase overrides `_apply` and adds weakrefs to params
-                   DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_to_empty', active_if=lambda p: p['swap']),),
+                   DecorateInfo(skipMPS),),
                decorators=rnn_gru_lstm_module_info_decorators),
     ModuleInfo(torch.nn.ReflectionPad1d,
                module_inputs_func=module_inputs_torch_nn_ReflectionPad1d,