[bug] workaround for idx fix

2025-09-09 21:09:18 +00:00 · 2024-05-28 02:33:12 +00:00
parent e0dde8fda5
commit 936dd96dbb
2 changed files with 7 additions and 8 deletions
--- a/colossalai/zero/gemini/gemini_mgr.py
+++ b/colossalai/zero/gemini/gemini_mgr.py
@@ -57,6 +57,7 @@ class GeminiManager:
        self._comp_cuda_demand_time = 0

    def reset_attributes(self):
+        assert self._compute_idx + 1 == len(self._compute_list)
        self._compute_idx = -1
        self._h2d_volume = 0
        self._d2h_volume = 0