From f9331ae2d09b066d330161737ec983804c169317 Mon Sep 17 00:00:00 2001 From: manfeibaigithub Date: Thu, 21 Mar 2024 17:38:20 +0000 Subject: [PATCH] test range --- torch_xla/experimental/fori_loop.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torch_xla/experimental/fori_loop.py b/torch_xla/experimental/fori_loop.py index 8a8e1f9c4e9..8986596607c 100644 --- 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))