examples/research_projects/jax-projects/big_bird/bigbird_flax.py/Args/__post_init__
class Args:    def __post_init__(self):
        os.makedirs(self.base_dir, exist_ok=True)
        self.save_dir = os.path.join(self.base_dir, self.save_dir)
        self.batch_size = self.batch_size_per_device * jax.device_count()