Solved by Manuel_Alejandro_Dia in post #2 
              
              
                I found the answer using torch.gather: 
residuals = torch.gather(residuals, 1, bin_logits.unsqueeze(1)).squeeze(1)

What I needed to do was to make the number of dimensions for the index array and the residuals equal. Then collapse the expanded dimension, et voilá!