Merge pull request #5738 from botbw/prefetch

[chore] fix init error
hpcaitech · May 21, 2024 · 83716e9 · 83716e9
2 parents f5b7de3 + b3c0e6d
commit 83716e9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/colossalai/zero/gemini/gemini_mgr.py b/colossalai/zero/gemini/gemini_mgr.py
@@ -42,7 +42,9 @@ def __init__(
         self._mem_stats_collector = (
             ChunkMemStatsCollector(chunk_manager, self._memstats) if policy_cls.need_mem_stats else None
         )
-        self._placement_policy = policy_cls(self, chunk_manager, self._mem_stats_collector, **placement_kwargs)
+        self._placement_policy = policy_cls(
+            chunk_manager=chunk_manager, mem_stats_collector=self._mem_stats_collector, **placement_kwargs
+        )
         self._compute_list: List[Tuple[Chunk, ...]] = []
         self._compute_idx: int = -1
         self._async_works: Dict[Chunk, dist.Work] = {}