@ptrblck thanks for the reply. It worked fine. How can I apply this where condition along y-axis? I have mask of shape torch.Size([2048]) and tensor of shape torch.Size([62, 2048]). I want to apply the mask along y-axis, can you please comment? Thanks.