fixed bernoulli1 and bernoulli2 tests

zero000064 · pytorchmergebot · commit 97001e1d6c56 · 2025-04-07T01:37:35.000Z
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -7944,9 +7944,7 @@ def test_bernoulli1(self):
         def fn(a):
             b = a.clone()
             # aten.bernoulli_() uses aten.bernoulli.p() behind the scene, so it will be decomposed.
-            return aten.bernoulli_(b).sum() / torch.prod(
-                torch.tensor(a.size(), device=a.device)
-            )
+            return aten.bernoulli_(b).sum() / torch.prod(torch.tensor(a.size()))
 
         p = 0.3
         self.common(
@@ -7961,9 +7959,7 @@ def fn(a):
     @skip_if_triton_cpu
     def test_bernoulli2(self):
         def fn(a):
-            return aten.bernoulli(a).sum() / torch.prod(
-                torch.tensor(a.size(), device=a.device)
-            )
+            return aten.bernoulli(a).sum() / torch.prod(torch.tensor(a.size()))
 
         p = 0.3
         self.common(
