Solved by KFrank in post #2 
              
              
                Hello Yoonho - 
Let me try to answer part of your question. 
(I haven’t looked at your dataloader or model code, so I won’t comment 
on it.) 
Let me try to give you some general context here. 

The shape of your input doesn’t really matter to your loss function – it 
just has to match what your mode…