@ptrblck
Thanks, it works now. But why was dim1 causing problems? Also I was unable to use the CrossEntropyLoss. It gave the error:
 RuntimeError: multi-target not supported at /pytorch/aten/src/THNN/generic/ClassNLLCriterion.c:21

Can’t I treat the first three inputs as of class 0 and the 4th input as of class 1?