I’m learning PyTorch by writing a simple neural network with one hidden layer. The dataset is from Mnist. I’m reading the data from a csv file, taken from Kaggle. This is my code.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

import pandas as pd
import numpy as np

import sys


class SingleHidden(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(SingleHidden, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.out = nn.Linear(H, D_out)
    
    def forward(self, x):
        print(x.size())
        print(self.linear1.weight.size())
        x = self.linear1(x)
        hidden_relu = nn.functional.relu(x)
        y_pred = self.out(x)
        return y_pred

class ReadData(Dataset):
    def __init__(self, file_name):
        self.dataframe = pd.read_csv(file_name)
    
    def __len__(self):
        return len(self.dataframe)
    
    def __getitem__(self, index):
        val = self.dataframe.loc[index]
        val = val.values.reshape(val.shape[0], 1)
        label = np.asarray(val[0]) #, dtype=np.float32)
        features = np.asarray(val[1:]) #, dtype=np.float32)
        data_obj = {
            'label': label,
            'features': features
        }
        return data_obj


N, D_in, H, D_out = 4, 784, 100, 10

if len(sys.argv) != 2:
    print('Pass training csv file for mnist')
    sys.exit(0)

dataset = ReadData(sys.argv[1])
dataloader = DataLoader(dataset=dataset, batch_size=4)

_brain = SingleHidden(D_in=D_in, H=H, D_out=D_out)
criterion = nn.CrossEntropyLoss()
# criterion = nn.MSELoss()
optimizer = optim.SGD(_brain.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):
    for obj in dataloader:
        features = obj['features']
        labels = obj['label']
        print(labels)
        optimizer.zero_grad()

        features = features.float()
        labels = labels.float()
        print(features.shape)

        output = _brain(features)
        loss = criterion(output, labels)

        # Compute gradients for all backprop elements.
        loss.backward()

        optimizer.step()

print('Finished')

The output that I get is
tensor([[1],
        [0],
        [1],
        [4]])
torch.Size([4, 784, 1])
torch.Size([4, 784, 1])
torch.Size([100, 784])
Traceback (most recent call last):
  File "SingleHidden.py", line 70, in <module>
    output = _brain(features)
  File "/home/sri/anaconda3/envs/ptorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 477, in __call__
    result = self.forward(*input, **kwargs)
  File "SingleHidden.py", line 21, in forward
    x = self.linear1(x)
  File "/home/sri/anaconda3/envs/ptorch/lib/python3.6/site-packages/torch/nn/modules/module.py", line 477, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/sri/anaconda3/envs/ptorch/lib/python3.6/site-packages/torch/nn/modules/linear.py", line 55, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/sri/anaconda3/envs/ptorch/lib/python3.6/site-packages/torch/nn/functional.py", line 1026, in linear
    output = input.matmul(weight.t())
RuntimeError: size mismatch, m1: [3136 x 1], m2: [784 x 100] at /opt/conda/conda-bld/pytorch_1532579245307/work/aten/src/TH/generic/THTensorMath.cpp:2070

What is causing this error?