fix condition to extend cache

microsoft · May 1, 2024 · 822aeee · 822aeee
1 parent b3c0750
commit 822aeee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
@@ -1321,7 +1321,7 @@ def _allocate_buffers(self, shapes_and_dtypes, requires_grad=False, num_buffers=
 
     def _allocate_or_extend_buffers(self, idx, shape, dtype):
         numel = reduce(mul, shape) if len(shape) > 0 else 1
-        if len(self._grad_layer_buf) <= idx or self._grad_layer_buf[idx].numel() > numel:
+        if len(self._grad_layer_buf) <= idx or self._grad_layer_buf[idx].numel() < numel:
             new_buf = self._allocate_buffer(shape, dtype=dtype, num_buffers=1)[0]
             if len(self._grad_layer_buf) <= idx:
                 self._grad_layer_buf.append(new_buf)