diff --git a/torch_xla/experimental/fori_loop.py b/torch_xla/experimental/fori_loop.py index fbdc94ce339..8946acbebe6 100644 --- a/torch_xla/experimental/fori_loop.py +++ b/torch_xla/experimental/fori_loop.py @@ -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, 1)).clone(), (body_fun(*init_vals)).clone(), init_vals[1].clone()) # init_vals[1:]) + return (upper.clone(), (torch.add(lower.clone(), 1)).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))