Only visualize XLA tensors

Author: tengyifei

torchprime/torch_xla_models/train.py

@@ -204,13 +204,14 @@ def _log_shapes(self, batch):
         )
         logger.info(f"[{self.name}] data shapes: {shapes}")
 
-        # Visualize one tensor.
-        import click
-        from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
-
+        # Visualize one example tensor.
         t = next(iter(pytree.tree_iter(batch)))
-        generated_table = visualize_tensor_sharding(t, use_color=False)
-        click.echo(generated_table)
+        if t.device.type == "xla":
+          import click
+          from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
+
+          generated_table = visualize_tensor_sharding(t, use_color=False)
+          click.echo(generated_table)
 
       def __len__(self):
         return len(self.dataloader)