g# This is a combination of 2 commits.

[chat] fix gemini strategy fox
hpcaitech · Sep 13, 2023 · 81c662a · 81c662a
1 parent 23281b1
commit 81c662a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/applications/Chat/coati/trainer/strategies/colossalai.py b/applications/Chat/coati/trainer/strategies/colossalai.py
@@ -156,7 +156,7 @@ def __init__(
         warnings.warn(f'Stage 3 only supports fp16. Precision is set to fp16.')
 
         # NOTE: dist should be initialized before calling get_current_device()
-        plugin_initializer = lambda: GeminiPlugin(chunk_init_device=get_current_device(),
+        plugin_initializer = lambda: GeminiPlugin(device=get_current_device(),
                                                   placement_policy=placement_policy,
                                                   precision='fp16',
                                                   pin_memory=pin_memory,