From 2a0b9412bf61efe2324b686559b31cec667e82a5 Mon Sep 17 00:00:00 2001
From: Manfei <41607353+ManfeiBai@users.noreply.github.com>
Date: Mon, 27 Nov 2023 15:20:19 -0800
Subject: [PATCH] Update __init__.py

---
 torch_xla/distributed/spmd/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torch_xla/distributed/spmd/__init__.py b/torch_xla/distributed/spmd/__init__.py
index fe661364842..39b1e16c03d 100644
--- 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",