Solved by Michael_Christiansen in post #2 
              
              
                I was seeing this because I was missing a dimension in my tensor, I suspect that’s what you’re seeing too. When you pass your full dataset the shape of x is torch.Size([50000, , , ]). When you index it down to one sample it’s shape is torch.Size([, , ]), but you want it to be torch.Size([1, , , ]). 
…