runtime error

jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these. The above exception was the direct cause of the following exception: Traceback (most recent call last): File "/home/user/app/app.py", line 9, in <module> model = FlaxBigBirdForQuestionAnswering.from_pretrained(FLAX_MODEL_ID, block_size=64, num_random_blocks=3) File "/usr/local/lib/python3.10/site-packages/transformers/modeling_flax_utils.py", line 901, in from_pretrained model = cls(config, *model_args, _do_init=_do_init, **model_kwargs) File "/usr/local/lib/python3.10/site-packages/transformers/models/big_bird/modeling_flax_big_bird.py", line 1640, in __init__ super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init) File "/usr/local/lib/python3.10/site-packages/transformers/modeling_flax_utils.py", line 220, in __init__ random_params = self.init_weights(self.key, input_shape) File "/usr/local/lib/python3.10/site-packages/transformers/models/big_bird/modeling_flax_big_bird.py", line 1676, in init_weights module_init_outputs = self.module.init( File "/usr/local/lib/python3.10/site-packages/transformers/models/big_bird/modeling_flax_big_bird.py", line 2410, in __call__ start_logits, end_logits = logits.split(self.config.num_labels, axis=-1) AttributeError: 'jaxlib.xla_extension.ArrayImpl' object has no attribute 'split'. Did you mean: '_split'?

Container logs: