add wandb integration

README.md

@@ -1 +1,6 @@
+---
+widget:
+- text: "Sewindu sudah kita tak berjumpa, rinduku padamu sudah tak terkira."
+---
+
 # GPT2-medium-indonesian

run_clm_flax.py

@@ -606,8 +606,6 @@ def main():
                 epochs.write(
                     f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"
                 )
-                if jax.process_index() == 0:
-                    wandb.log({'my_metric': train_metrics})
 
                 train_metrics = []
 
@@ -640,7 +638,6 @@ def main():
                 if has_tensorboard and jax.process_index() == 0:
                     cur_step = epoch * (len(train_dataset) // train_batch_size)
                     write_eval_metric(summary_writer, eval_metrics, cur_step)
-                    wandb.log({'my_metric': eval_metrics})
 
             if cur_step % training_args.save_steps == 0 and cur_step > 0:
                 # save checkpoint after each epoch and push checkpoint to the hub