small fix

stanford-crfm · Dec 15, 2024 · 07781e6 · 07781e6
1 parent f7f2382
commit 07781e6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/levanter/optim/kron.py b/src/levanter/optim/kron.py
@@ -507,7 +507,7 @@ def broadcast_qs(_, ps, q, s):
 
         if return_partition_specs_only:
             return dict(
-                key=jax.random.PRNGKey(jax.process_index()),
+                key=PartitionSpec(),
                 count=PartitionSpec(),
                 mu=mu_sharding,
                 Qs_preconditioners=Qs_sharding,
@@ -516,6 +516,7 @@ def broadcast_qs(_, ps, q, s):
             )
 
         return dict(
+            key=jax.random.PRNGKey(jax.process_index()),
             count=jnp.zeros([], jnp.int32),
             mu=mu,
             Qs_preconditioners=Qs,