Revert "Default XLA to use swap_tensors path in nn.Module._apply (#126814)" (#128170)

mikaylagawarecki · pytorchmergebot · commit 65aa16f968af · 2024-06-07T01:44:14.000Z
#128165 :( This reverts commit a7b1dd8. Pull Request resolved: #128170 Approved by: https://github.com/drisspg, https://github.com/albanD
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -8184,9 +8184,9 @@ def test_batchnorm_large_batch(self, device, dtype):
     @dtypes(torch.float, torch.double, torch.bfloat16, torch.complex128)
     def test_conv_empty_input(self, device, dtype):
         def help(input, conv, memory_format):
-            ref_out = conv(input).detach()
+            ref_out = conv(input)
             conv_cl = conv.to(memory_format=memory_format)
-            out_cl = conv_cl(input).detach()
+            out_cl = conv_cl(input)
             self.assertEqual(ref_out, out_cl)
             input_cl = input.to(memory_format=memory_format)
             out_cl2 = conv(input_cl)
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
@@ -794,13 +794,6 @@ def compute_should_use_set_data(tensor, tensor_applied):
 
         should_use_swap_tensors = torch.__future__.get_swap_module_params_on_conversion()
 
-        def compute_should_use_swap_tensors(tensor, tensor_applied):
-            return (should_use_swap_tensors
-                    # subclasses may have multiple child tensors so we need to use swap_tensors
-                    or is_traceable_wrapper_subclass(tensor_applied)
-                    or tensor.device.type == 'xla'
-                    or tensor_applied.device.type == 'xla')
-
         for key, param in self._parameters.items():
             if param is None:
                 continue
@@ -811,7 +804,8 @@ def compute_should_use_swap_tensors(tensor, tensor_applied):
                 param_applied = fn(param)
             p_should_use_set_data = compute_should_use_set_data(param, param_applied)
 
-            p_should_use_swap_tensors = compute_should_use_swap_tensors(param, param_applied)
+            # subclasses may have multiple child tensors so we need to use swap_tensors
+            p_should_use_swap_tensors = should_use_swap_tensors or is_traceable_wrapper_subclass(param_applied)
 
             param_grad = param.grad
             if p_should_use_swap_tensors:
