updated 50 train files

src/train.py

@@ -206,7 +206,7 @@ def main():
         if idx%5==0:
           summary_writer.scalar('train_loss',flax.jax_utils.unreplicate(train_metric)['loss'].item(),idx)
           summary_writer.scalar('train_accuracy', jax.device_get(train_metric['accuracy']).mean().item(),idx)
-        if idx%20==0:
+        if idx%50==0:
             logger.info(f"train_step_loss{idx}:   {flax.jax_utils.unreplicate(train_metric)['loss'].item()} train_step_acc{idx}:   {jax.device_get(train_metric['accuracy']).mean().item()}")
             
         progress_bar_train.update(1)