U Net won't train. Any ideas?

I know I am being a bit lazy here by asking this question.

I am trying to train a UNet model on Kvasir medical image dataset. My dice loss doesn’t decrease.

I included a basic block to check if the weights of the layers were being updated.

It tunrs out that only my final hidden layer is being updated. Which means the training might have vanishing gradient problem.

Ps- I am trying to train the network without BatchNorm layers.