test range

pytorch · Mar 21, 2024 · 1e449b5 · 1e449b5
1 parent 6bf4f4e
commit 1e449b5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
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
@@ -87,7 +87,7 @@ def body_fun(a, b):
       return torch.add(a, b) # [0])
     # _, _, res, _ = fori_loop(lower, upper, body_fun, init_val, one_value) # init_val_list) # init_val)
     # A, B, res, D = fori_loop(lower, upper, body_fun, init_val, one_value) # init_val_list) # init_val)
-    A, B, res, D = fori_loop(upper, lowers) # lower, upper, body_fun, init_val, one_value)
+    A, B, res, D = fori_loop(upper, body_fun, lowers) # lower, upper, body_fun, init_val, one_value)
     print("result: ", res) # init_val_
     print("A: ", A) # lower_
     print("B: ", B) # upper_

diff --git a/torch_xla/experimental/fori_loop.py b/torch_xla/experimental/fori_loop.py
@@ -12,7 +12,7 @@
 from torch._higher_order_ops.while_loop import while_loop_op
 
 
-def fori_loop(upper, lowers):#  upper, body_fun, *init_vals): # *init_val):
+def fori_loop(upper, body_fun, lowers):#  upper, body_fun, *init_vals): # *init_val):
 
   device = xm.xla_device()
   # limit_value = upper