Skip to content

Commit

Permalink
test range
Browse files Browse the repository at this point in the history
  • Loading branch information
ManfeiBai committed Mar 21, 2024
1 parent ec97684 commit 9bcb354
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion torch_xla/experimental/fori_loop.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def body_fn(upper, lower, *init_vals):
# one_value_original = torch.tensor(1, dtype=torch.int32, device=device)
# (a, b) = init_vals
# return (upper, torch.add(lower, 1), body_fun(a, b), b.clone())
return (upper.clone(), (torch.add(lower.clone(), 1)).clone(), (body_fun(*init_vals)).clone(), init_vals[1].clone()) # init_vals[1:])
return (upper.clone(), (torch.add(lower.clone(), init_vals[1].clone())).clone(), (body_fun(*init_vals)).clone(), init_vals[1].clone()) # init_vals[1:])
# (body_fun(*init_vals)).clone(), init_vals[1].clone())
# body_fun(one_value_original, init_val)) # body_fun(lower, init_val))

Expand Down

0 comments on commit 9bcb354

Please sign in to comment.