Validation loss doesn't decrease when using Lightning

HI everyone! I’m trying to build a segmentation model with segmentation models pytorch and Lightning. However, when I try to fit my model I see that the validation loss is always pretty high (even if I plot it in tensorboard, it doesn’t seem to decrease at all, but it just goes up and down). The same doesn’t happen when I train the same model in pure Pytorch. The training loss instead seems to keep decreasing as expected.

You can see my Lightning implementation and try to run the code yourself this Colab link.