[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
2025-09-10 05:20:33 +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/common.py
+++ b/colossalai/utils/common.py
@@ -10,6 +10,15 @@ from typing import Callable
 import numpy as np
 import torch

+from colossalai.accelerator import get_accelerator
+
+
+def get_current_device():
+    """
+    A wrapper function for accelerator's API for backward compatibility.
+    """
+    return get_accelerator().get_current_device()
+

 def ensure_path_exists(filename: str):
    # ensure the path exists