Set correct device id on efficientzerotensors (#71611)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71611
Fixes https://github.com/pytorch/pytorch/issues/71160 https://github.com/pytorch/pytorch/issues/69925 #69913
Test Plan: Imported from OSS
Reviewed By: VitalyFedyunin
Differential Revision: D33897543
Pulled By: anjali411
fbshipit-source-id: f1d8608c351876b8c2619da5ef891f74bad30ab5
(cherry picked from commit 643e666ea34aaf1584d478cd85025bf66d93cb21)
diff --git a/test/test_torch.py b/test/test_torch.py
index 4a5ac3b..164e658 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -5477,6 +5477,10 @@
actual = x[..., :1].clamp(lb, ub)
self.assertEqual(expect, actual)
+ def test_cuda_device_idx(self, device):
+ x = torch.zeros(3, device=device)
+ y = torch._efficientzerotensor(3, device=device)
+ self.assertEqual(x.device, y.device)
# we implemented custom deallocation for subclasses, so it behooves
# us to make sure all of these bits work. We'll use __del__ to
