test range

pytorch · Mar 21, 2024 · 80d28c4 · 80d28c4
1 parent a197c8a
commit 80d28c4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torch_xla/experimental/fori_loop.py b/torch_xla/experimental/fori_loop.py
@@ -84,6 +84,9 @@ def _xla_while_loop(cond_fn, body_fn, operands):
   cond_hlo = cond_ctx.hlo()
   cond_computation = xb.computation_from_module_proto("condcomputation",
                                                       cond_hlo)
+  cond_hlo_print = xb.get_computation_hlo(cond_computation)
+  print("cond computation: !!!!!!!!!")
+  print(cond_hlo_print)
 
   # generate body_fn xlacomputation
   body_result = body_fn(*operands)