[ColoTensor] rename APIs and add output_replicate to ComputeSpec (#1168)

2025-09-07 12:01:39 +00:00 · 2022-06-24 13:08:54 +08:00
parent f4ef224358
commit 4b9bba8116
23 changed files with 116 additions and 105 deletions
--- a/colossalai/tensor/param_op_hook.py
+++ b/colossalai/tensor/param_op_hook.py
@@ -129,7 +129,7 @@ def _get_colo_tensors_info(*args) -> list:
    info = []
    for arg in args:
        if isinstance(arg, ColoTensor):
-            info.append((arg.__class__, arg.spec))
+            info.append((arg.__class__, arg.tensor_spec))
        else:
            info.append(None)
    return info