Solved by Algomorph in post #2 
              
              
                Big thanks to Ahmed Taha, who (in another thread / forum) suggested to replace the 
mask[mask > 0.999] = 1.0 
with 
mask = torch.where(mask > 0.999, torch.ones_like(mask), mask) 
This solution didn’t produce any errors so far, so I think it is valid. Why the original code produced an error for some …