Update __init__.py

pytorch · Nov 27, 2023 · 2a0b941 · 2a0b941
1 parent 14e2bed
commit 2a0b941
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_xla/distributed/spmd/__init__.py b/torch_xla/distributed/spmd/__init__.py
@@ -3,7 +3,7 @@
                            XLAPatchedLinear, mark_sharding, clear_sharding,
                            wrap_if_sharded, xla_patched_nn_linear_forward)
 from .api import xla_distribute_tensor, xla_distribute_module
-# from .debugging import visualize_tensor_sharding
+from .debugging import visualize_tensor_sharding
 
 __all__ = [
     "XLAShard",