Hey The size of my images is as follows
TrainImages = torch.FloatTensor(1200,3,256, 256)
TrainLabels = torch.LongTensor(1200)
TestImages = torch.FloatTensor(400,3,256,256)
TestLabels = torch.LongTensor(400)

Hope this helps!
edit I tried printing it out and it gave the following output
torch.Size([1200, 3, 256, 256])
torch.Size([1200])
torch.Size([400, 3, 256, 256])
torch.Size([400])