fix: safer pipeline.clear()

instadeepai · Oct 11, 2024 · 4748636 · 4748636
1 parent 8be8037
commit 4748636
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mava/utils/sebulba.py b/mava/utils/sebulba.py
@@ -142,7 +142,10 @@ def get(
     def clear(self) -> None:
         """Clear the pipeline."""
         while not self._queue.empty():
-            self._queue.get()
+            try:
+                self._queue.get(block=False)
+            except queue.Empty:
+                break
 
     def shard_split_playload(self, payload: Any, axis: int = 0) -> Any:
         split_payload = jnp.split(payload, len(self.learner_devices), axis=axis)