[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
2025-09-12 12:47:21 +00:00 · 2024-01-25 17:01:48 +08:00
parent d7f8db8e21
commit 7cfed5f076
157 changed files with 1353 additions and 8966 deletions
--- a/colossalai/utils/init.py
+++ b/colossalai/utils/init.py
@@ -4,6 +4,7 @@ from .common import (
    disposable,
    ensure_path_exists,
    free_storage,
+    get_current_device,
    is_ddp_ignored,
    set_seed,
 )
@@ -22,5 +23,6 @@ __all__ = [
    "_cast_float",
    "free_storage",
    "set_seed",
+    "get_current_device",
    "is_ddp_ignored",
 ]