Binary Segmentation project code review

I am developing a binary segmentation model based on a ViT (Vision Transformer) backbone and a binary segmentation mlp "decoder" following [1]. I have the whole pipeline already implemented for the TUSimple dataset [2] but I am struggling to even get F1 or IoU scores above 0.01. I experimented with Dice Loss or BCEwithLogits functions. I need somebody to conduct code review and identify what could be wrong as I believe there is something faulty during training causing gradients to explode. Code for review along with more information on the project will be provided at request.


[1] https://arxiv.org/abs/2105.05633
[2] https://github.com/TuSimple/tusimple-benchmark/issues/3