Thanks everyone for the help! following you advices I managed to modify and run correctly the unet (more or less, the loss kind of saturates and does not go particularly down but that might be my dataset).
Quick question is there a way to easely implement dropout in the U-net snippet found here https://github.com/usuyama/pytorch-unet ?
Best and thanks again for the help!