Training Error on the Tutorial (FCN8 and UNET Semantic Segmentation with Keras and Xilinx Vitis AI)

[HA-999]

I am working on the following tutorial.
https://github.com/Xilinx/Vitis-AI-Tutorials/blob/1.4/Design_Tutorials/05-Keras_FCN8_UNET_segmentation/README.md
My goal is to train this segmentation model using a custom dataset and run it on Ultra96.
I trained the model using a sample dataset on a CPU docker, however, I could not train it on a GPU docker properly (no increase in mAP).
I guess the wrong configuration of the GPU docker cause the problem however, I cannot find the solution. How can I train the model properly using a GPU?
Thank you.

[codimahi]

I am also having the same issue.
Were you able to solve it?