small fixes

ml-energy · Dec 13, 2024 · ea0f866 · ea0f866
1 parent 63f7fe9
commit ea0f866
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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