Semi supervised model not converging with only Dice loss

[nicolas1805961]

Hello.

I am trying to train the semi supervised model on the cardiac ACDC dataset. As I want to obtain the best Dice score as possible, I have set the weight of each loss to zero except the Dice loss which has a weight of 1 (loss with labels).
The problem is that the model does not converge. After a few epochs, there seems to be a collapse and the Dice loss suddenly jumps. The more parameters the model has, the earlier in training this collapse happens (the collapse is almost immediate, before the first epoch with ~20 M parameters but appears after 200 epochs when having only 100 000 parameters). When the collapse happens, the "vxm_dense_flow_loss" or "vxm_dense_flow_resize_loss" (these losses have a 0 weight) explode to a very high value (millions, sometimes a little bit less).

When registering with the learnt weights, the output is an all black image.

Here is what the flow looks like after 600 epochs for the x axis for the smallest model which collapses after 200 epochs:

Values of the flow are between 0 and -250.

Input image values are normalized between 0 and 1.
I have tried to increase the batch size from 1 to 16. I also tried to reduce the learning rate from 1e-4 to 1e-5. It did not solve the issue.

Would you have any idea as to why this happens ?

Thank you for your help.