Thanks for the stack trace and code.
Based on the code and error message, I would guess outputs[x][mask] or disp_L[mask] is raising this error.
Could you check the shape of all tensors in this operation and the values of mask?