So it’s not the same putting the unsqueesze in the dataset class that after that.
Dataloader cat several samples to form a batch.
So if you put it in the dataset class you are stucking 2 tensors of shape 1x2048.
If you unsqueeze once the dataloader returns the samples you get samples of 2x2048. Then it should be unsqueeze(1)