diff --git a/distributed_shampoo/tests/shampoo_model_utils_test.py b/distributed_shampoo/tests/shampoo_model_utils_test.py
index 88c9afd..bf477b7 100644
--- a/distributed_shampoo/tests/shampoo_model_utils_test.py
+++ b/distributed_shampoo/tests/shampoo_model_utils_test.py
@@ -135,7 +135,3 @@ def test_initialization(self):
                     torch.testing.assert_close(
                         original_linear.weight, combined_linear.combined_weight
                     )
-
-
-if __name__ == "__main__":
-    unittest.main()