slight update

src/gptneo_story.py

@@ -123,7 +123,8 @@ def train_step(state,batch,dropout_rng):
         #Added.
   logits=new_state.apply_fn(**batch,params=new_state.params,dropout_rng=dropout_rng,train=True)[0]
   accuracy=jnp.equal(jnp.argmax(logits,axis=-1),targets)
-  metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
+  #metrics=jax.lax.pmean({"loss":loss,"learning_rate":learning_rate_function(state.step),'accuracy':accuracy},axis_name="batch")
+  metrics=jax.lax.pmean({"loss":jax.device_get(loss),"learning_rate":jax.device_get(learning_rate_function(state.step)),'accuracy':jax.device_get(accuracy)},axis_name="batch")
   return new_state,metrics,new_dropout_rng
 
 parallel_train_step = jax.pmap(train_step, axis_name="batch", donate_argnums=(0,))
@@ -135,7 +136,8 @@ def eval_step(state, batch):
   predictions=state.logits_function(logits)
   eval_accuracy=jnp.equal(predictions,targets)
   #eval_acc=jnp.equal(predictions,targets)
-  metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
+  metrics=jax.lax.pmean({"loss":jax.device_get(loss),'accuracy':jax.device_get(eval_accuracy)},axis_name="batch")
+  #metrics=jax.lax.pmean({"loss":loss,'accuracy':eval_accuracy},axis_name="batch")
   #return state.logits_function(logits)  #(8,4)
   return targets,predictions,metrics