Could you explain what these lines of code are doing and what values are stored in segments afterwards?
    segments = segments*255
    segments = segments.int().long()
    segments[segments==255] = 21

Based on your code you might be normalizing the targets and the mentioned lines of code might try to denormalize them.
If you’ve checked the segments and made sure thay contain the right values, I would recommend to try to overfit a small data sample (e.g. just 10 samples) and make sure your mode is able to do so.