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 cases remains a mystery to me.