[plugin] hotfix zero plugin (#6036)

* [plugin] hotfix zero plugin * [plugin] hotfix zero plugin
hpcaitech · Aug 28, 2024 · cc1b0ef · cc1b0ef
1 parent d383449
commit cc1b0ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/colossalai/booster/plugin/low_level_zero_plugin.py b/colossalai/booster/plugin/low_level_zero_plugin.py
@@ -100,7 +100,7 @@ def forward(self, *args, **kwargs):
         if self.convert_fn is not None:
             args = tree_map(self.convert_fn, args)
             kwargs = tree_map(self.convert_fn, kwargs)
-        ctx = ColoParamOpHookManager.use_hooks(self.op_hook) if self.overlap_allgather else nullcontext()
+        ctx = ColoParamOpHookManager.use_hooks(*self.op_hooks) if self.overlap_allgather else nullcontext()
         with ctx:
             return super().forward(*args, **kwargs)