Alternatively to the dice loss you could also use e.g. nn.CrossentropyLoss, if each pixel belongs to one certain class (note that you would need to pass the logits to this loss function).
The training loop looks alright. I’m wondering, why you are calling empty_cache() inside the loop, as this could make your code slower and shouldn’t save any memory.