I managed to make this work, I loaded the model and saved the dict and then I created a Unet with dropout in the architecture and loaded the dict.
However, I am getting this error:
RuntimeError: Error(s) in loading state_dict for Unet:
	Missing key(s) in state_dict: "segmentation_head.1.weight", "segmentation_head.1.bias". 
	Unexpected key(s) in state_dict: "segmentation_head.0.weight", "segmentation_head.0.bias".

I believe this is because of the modification I did in the SegmentatioHead by adding dropout. I want to load the unexpected keys as they are the missing keys. Is there any way to do this?