examples/research_projects/jax-projects/big_bird/bigbird_flax.py/Trainer/train
class Trainer:    def train(self, state, tr_dataset, val_dataset):
        args = self.args
        total = len(tr_dataset) // args.batch_size

        rng = jax.random.PRNGKey(0)
        drp_rng = jax.random.split(rng, jax.device_count())
        for epoch in range(args.max_epochs):
            running_loss = jnp.array(0, dtype=jnp.float32)
            tr_dataloader = get_batched_dataset(tr_dataset, args.batch_size, seed=epoch)
            i = 0
            for batch in tqdm(tr_dataloader, total=total, desc=f"Running EPOCH-{epoch}"):
                batch = self.data_collator(batch)
                state, metrics, drp_rng = self.train_step_fn(state, drp_rng, **batch)
                running_loss += jax_utils.unreplicate(metrics["loss"])
                i += 1
                if i % args.logging_steps == 0:
                    state_step = jax_utils.unreplicate(state.step)
                    tr_loss = running_loss.item() / i
                    lr = self.scheduler_fn(state_step - 1)

                    eval_loss = self.evaluate(state, val_dataset)
                    logging_dict = {
                        "step": state_step.item(),
                        "eval_loss": eval_loss.item(),
                        "tr_loss": tr_loss,
                        "lr": lr.item(),
                    }
                    tqdm.write(str(logging_dict))
                    self.logger.log(logging_dict, commit=True)

                if i % args.save_steps == 0:
                    self.save_checkpoint(args.save_dir + f"-e{epoch}-s{i}", state=state)