So these are my parameters shape

print(model.parameters())
print(len(list(model.parameters())))
print(list(model.parameters())[0].size())
print(list(model.parameters())[1].size())

output:
<generator object Module.parameters at 0x7fc5eb5a6fc0>
2
torch.Size([2, 784])
torch.Size([2])

The shape it takes after using DataLoader:

for i, (images, labels) in enumerate(train_loader):
print(images.shape)

output:
This goes on for 36 times
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])
torch.Size([50, 3, 28, 28])


images = Variable(images.view(-1, 28*28).cuda())
image size after the above line
output:
torch.Size([150, 784])


outputs = model(images)
outputs size after the above line
output:
torch.Size([150, 2])