[zeta.utils][__init__][CLEANUP]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zetascale"
-version = "0.9.3"
+version = "0.9.4"
 description = "Transformers at zeta scales"
 authors = ["Zeta Team <kye@apac.ai>"]
 license = "MIT"

zeta/utils/__init__.py

@@ -1,3 +1,17 @@
 # Copyright (c) 2022 Agora
 # Licensed under The MIT License [see LICENSE for details]
-from zeta.utils.main import *
+from zeta.utils.cuda_memory_wrapper import track_cuda_memory_usage
+
+from zeta.utils.benchmark import (
+    benchmark,
+    print_cuda_memory_usage,
+    save_memory_snapshot,
+)
+
+
+__all__ = [
+    "track_cuda_memory_usage",
+    "benchmark",
+    "print_cuda_memory_usage",
+    "save_memory_snapshot",
+]