[builder] unified cpu_optim fused_optim inferface (#2190)

2025-09-12 12:47:21 +00:00 · 2022-12-23 20:57:41 +08:00
parent 9587b080ba
commit 355ffb386e
9 changed files with 28 additions and 50 deletions
--- a/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py
+++ b/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py
@@ -14,7 +14,7 @@ class MultiTensorApply(object):

    def __init__(self, chunk_size):
        try:
-            import colossalai._C.fused_optim
+            from colossalai.kernel import fused_optim
            MultiTensorApply.available = True
            self.chunk_size = chunk_size
        except ImportError as err: