Tensorboard error while running mlm_flax TPU example script on TPU
## Environment info
Following the https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-setup-tpu-vm

## Script

./run_mlm_flax.py \
    --output_dir="./" \
    --model_type="roberta" \
    --config_name="./" \
    --tokenizer_name="./" \
    --dataset_name="oscar" \
    --dataset_config_name="unshuffled_deduplicated_als" \
    --max_seq_length="128" \
    --per_device_train_batch_size="4" \
    --per_device_eval_batch_size="4" \
    --learning_rate="3e-4" \
    --warmup_steps="1000" \
    --overwrite_output_dir \
    --num_train_epochs="8" \
    --push_to_hub

## Error
v  File "./run_mlm_flax.py", line 66
    print(f"Unable to display metrics through TensorBoard because some package are not installed: {ie}")
                                                                                                      ^
## Information
Tensorboard is already installed using pip install tensorboard

Can you send me your tpu_name & name of virtual environment in a personal Slack message? I'll have a look<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.