Rejection sampling broken

[gtangg12]

Hi, I believe

chosen_indices_validity = mask.squeeze()[c, y, x].bool()

does not work for batch size 1 e.g. evaluation sampling? Currently, my eval mask has shape torch.Size([1, 640, 480, 1]), if that's correct (since I am using a custom datamanger).

The other methods seem to use

mask[..., 0]