test range

pytorch · Mar 21, 2024 · a197c8a · a197c8a
1 parent 1e449b5
commit a197c8a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py b/test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py
@@ -81,7 +81,7 @@ def test_fori_loop_tpu_addition(self):
     upper = torch.tensor([30], dtype=torch.int32, device=device)
     one_value = torch.tensor([1], dtype=torch.int32, device=device)
     init_val_list = (init_val, one_value)
-    lowers = torch.tensor([1, 1, 1], dtype=torch.int32, device=device) # lower, init_val, one_value
+    lowers = torch.tensor([[1], [1], [1]], dtype=torch.int32, device=device) # lower, init_val, one_value
 
     def body_fun(a, b):
       return torch.add(a, b) # [0])