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
index 1f8efbd482f..15eab77a34c 100644
--- 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
@@ -92,7 +92,7 @@ def test_fori_loop_tpu_simple_linear(self):
     device = xm.xla_device()
     torch.set_grad_enabled(False)
 
-    upper = torch.tensor([1], dtype=torch.int32, device=device)
+    upper = torch.tensor([50], dtype=torch.int32, device=device)
     lower = torch.tensor([0], dtype=torch.int32, device=device)
     init_val = torch.tensor([1], dtype=torch.int32, device=device)
     l_in_0 = torch.randn(10, device=xm.xla_device())