[gemini] fix missing return (#5845)

2026-02-21 06:22:09 +00:00 · 2024-06-21 11:38:40 +08:00
parent bd3e34fef6
commit 8a5c86439a
1 changed files with 1 additions and 0 deletions
--- a/colossalai/zero/gemini/gemini_ddp.py
+++ b/colossalai/zero/gemini/gemini_ddp.py
@@ -450,6 +450,7 @@ class GeminiDDP(ModelWrapper):
                    chunk_manager.move_chunk(grad_chunk, grads_device[p], force_copy=True)
                    if not (master_weights) or (enable_gradient_accumulation):
                        chunk_manager.move_chunk(chunk, grads_device[p], force_copy=True)
+        return empty_grad

    def zero_grad(self, set_to_none: bool = False) -> None:
        self.module.zero_grad(set_to_none=True)