System Info

• transformers version: 4.22.0
• Platform: Linux-4.19.157-1.20201118.el7.x86_64-x86_64-with-glibc2.17
• Python version: 3.8.3
• Huggingface_hub version: 0.9.1
• PyTorch version (GPU?): 1.10.2+cu111 (True)
• Tensorflow version (GPU?): 2.10.0 (True)
• Flax version (CPU?/GPU?/TPU?): 0.6.0 (gpu)
• Jax version: 0.3.17
• JaxLib version: 0.3.15
• Using GPU in script?: YES
• Using distributed or parallel set-up in script?: NO

GPU device: A100 x 1

Who can help?

@SaulLu

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

I run the following code
https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_mlm_flax.py#L296
and an error occurs here:

My required_input is a tensor (in jax), sometimes I get error ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all().

Expected behavior

I suggest to use if required_input is None instead of if not required_input to avoid error, as the latter one is ambiguous.