enable bf16 on eval (#46)

Author: jugheadjones10

training/train_single_task.py

@@ -260,8 +260,10 @@ def _update_minbatch(train_state, batch_info):
                 env,
                 env_params,
                 train_state,
-                # TODO: make this as a static method mb?
-                jnp.zeros((1, config.rnn_num_layers, config.rnn_hidden_dim)),
+                jnp.zeros(
+                    (1, config.rnn_num_layers, config.rnn_hidden_dim),
+                    dtype=jnp.bfloat16 if config.enable_bf16 else None,
+                ),
                 1,
             )
             eval_stats = jax.lax.pmean(eval_stats, axis_name="devices")