From ea0f8665e3efd04551e2516156e7148697f78168 Mon Sep 17 00:00:00 2001 From: Parth Raut Date: Thu, 12 Dec 2024 22:07:56 -0500 Subject: [PATCH] small fixes --- zeus/utils/framework.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/zeus/utils/framework.py b/zeus/utils/framework.py index 1e5cd749..b321e0e0 100644 --- a/zeus/utils/framework.py +++ b/zeus/utils/framework.py @@ -142,7 +142,7 @@ def all_reduce( if jax.process_count() == 1: return object - raise NotImplementedError("JAX all-reduce not yet implemented") + raise NotImplementedError("JAX distributed all-reduce not yet implemented") raise RuntimeError("No framework is available.") @@ -155,4 +155,3 @@ def is_distributed() -> bool: if jax_is_available(): jax = MODULE_CACHE["jax"] return jax.process_count() > 1 - return False