Hi @albanD, @spanev and @tom,
The problem was fixed when converting the PyTorch tensor to numpy, doing this operation and then converting back to PyTorch tensor, as follows:
    mask2 = mask.clone().cpu().numpy()
    values_batch = value.squeeze(1).cpu().numpy()
    mask2[mask2 == 1] = values_batch * 0.5
    mask2= torch.from_numpy(mask2).to(device)   

Hope this issue could be fixed in future PyTorch releases. Thanks!