fix CI

Author: JackCaoG

torch_xla/core/dynamo_bridge.py

@@ -18,7 +18,7 @@
 import torch_xla.runtime as xr
 import torch_xla.utils.utils as xu
 
-debug = os.environ.get("TORCH_XLA_DEBUG") == "1"
+debug = os.environ.get("XLA_DYNAMO_DEBUG") == "1"
 
 
 @dataclasses.dataclass
@@ -322,7 +322,12 @@ def extract_graph_helper(xla_model: torch.fx.GraphModule):
 
 
 def extract_internal(xla_model: torch.fx.GraphModule):
-  xm.mark_step()
+  if debug:
+    print('after partitioner')
+    xla_model._graph.print_tabular()
+  if any(
+      torch_xla._XLAC._check_tensor_need_materialization(xla_model.xla_args)):
+    xm.mark_step(wait=True)
   (xla_args_sharding_spec, args_and_out, graph_hash,
    arg_index_to_need_update_index, none_remover, graph_input_matcher,
    dumb_return_handler, xla_args_need_update) = extract_graph_helper(xla_model)
@@ -445,6 +450,8 @@ def call_module(self, target, args, kwargs):
 
 
 def extract_compiled_graph(xla_model: torch.fx.GraphModule, xla_args):
+  if debug:
+    xla_model._graph.print_tabular()
   # This call is critical to make sure xla_args' tensor id show up in graph_input_tensor_ids
   xm.mark_step()