Hi @Fawaz_Sammani,
What’s the PyTorch version? Did you try with another one?
What are the mask2 values when it fails? Is there (at least) a 1 in mask2?
Can you try also try with:
mask2 = mask.clone().detach()