RuntimeError: stack expects each tensor to be equal size, but got [3] at entry 0 and [2] at entry 1

1

I am getting an above error , I don’t have an idea why I am getting this error. @ptrblck sampler Module code Snippet for specific part where I am getting an error is below: