tmoroder

Hi @gante. > If I got it right, the second notebook replaces the ``take_along_axis function``, and the third notebook also replaces the custom dropout. Still, there are XLA exceptions. Correct....

> Hi @tmoroder, can you try on GPU with jit_compile=True in both 4.20 and 4.21? Using ``jit_compile=True`` while compiling the model gives an error for both 4.20.1 and 4.21, e.g.,...

@gante Thanks a lot for your effort. Maybe I am doing something wrong... but using the code from your pull request it now runs on GPU (with ``jit_compile=True`` as additional...

@gante Great, setting the ``batch_size`` works 🥳. I only had to make sure that it divides the ``strategy.num_replicas_in_sync``, [FineTuning_TF_DeBERTa_Working_Fix_TPU](https://colab.research.google.com/drive/1wQ_shM9zigRzeATvcncTC4koFb2GkDgY?usp=sharing). Thanks a lot, I will test the procedure now on my...

@WissamAntoun Here some output that I get during the ``model.fit`` call. The model is very close to the one in the Colab notebooks, but the run is carried out on...