Update test: cast arange to float to avoid divsion problem

Bluefog-Lib · Nov 13, 2020 · ce5df1a · ce5df1a
1 parent 59b98a6
commit ce5df1a
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/test/torch_win_ops_test.py b/test/torch_win_ops_test.py
@@ -340,7 +340,8 @@ def test_win_put_with_varied_tensor_elements(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
-            base_tensor = torch.arange(DIM_SIZE**dim).view_as(tensor)/1000
+            base_tensor = torch.arange(
+                DIM_SIZE**dim, dtype=torch.float32).view_as(tensor).div(1000)
             tensor = self.cast_and_place(tensor, dtype)
             base_tensor = self.cast_and_place(base_tensor, dtype)
             tensor = tensor + base_tensor
@@ -461,7 +462,8 @@ def test_win_accumulate_with_varied_tensor_elements(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
-            base_tensor = torch.arange(DIM_SIZE**dim).view_as(tensor)/1000
+            base_tensor = torch.arange(
+                DIM_SIZE**dim, dtype=torch.float32).view_as(tensor).div(1000)
             tensor = self.cast_and_place(tensor, dtype)
             base_tensor = self.cast_and_place(base_tensor, dtype)
             tensor = tensor + base_tensor
@@ -627,7 +629,8 @@ def test_win_get_with_varied_tensor_elements(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
-            base_tensor = torch.arange(DIM_SIZE**dim).view_as(tensor)/1000
+            base_tensor = torch.arange(
+                DIM_SIZE**dim, dtype=torch.float32).view_as(tensor).div(1000)
             tensor = self.cast_and_place(tensor, dtype)
             base_tensor = self.cast_and_place(base_tensor, dtype)
             tensor = tensor + base_tensor