test range

pytorch · Mar 21, 2024 · f9331ae · f9331ae
1 parent b32e2c2
commit f9331ae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_xla/experimental/fori_loop.py b/torch_xla/experimental/fori_loop.py
@@ -31,10 +31,10 @@ def cond_fn(upper, lower, *init_vals):
 
   def body_fn(upper, lower, *init_vals):
     # one_value_original = torch.tensor(1, dtype=torch.int32, device=device)
-    # (a, b) = init_vals
+    (a, b) = init_vals
     # return (upper, torch.add(lower, 1), body_fun(a, b), b.clone())
     # return (upper.clone(), (torch.add(lower.clone(), init_vals[1].clone())).clone(), (body_fun(*init_vals)).clone(), init_vals[1].clone()) # init_vals[1:])
-    return (upper, (torch.add(lower, init_vals[1])), (body_fun(*init_vals)), init_vals[1]) # init_vals[1:])
+    return (upper, (torch.add(lower, b)), (body_fun(*init_vals)), b) # init_vals[1:])
     # (body_fun(*init_vals)).clone(), init_vals[1].clone())
     # body_fun(one_value_original, init_val)) # body_fun(lower, init_val))