title: Two-Layer Neural Network in PyTorch does not Converge
body: <h1>Problem</h1>
  
  <p>I am trying to implement 2-layer neural network using different methods (TensorFlow, PyTorch and from scratch) and then compare their performance based on MNIST dataset.</p>
  
  <p>I am not sure what mistakes I have made, but the accuracy in PyTorch is only about 10%, which is basically random guess. I think probably the weights does not get updated at all. </p>
  
  <p>Note that I intentionally use the dataset provided by TensorFlow to keep the data I use through 3 different methods consistent for accurate comparison. </p>
  
  <pre class="lang-py prettyprint-override"><code>from tensorflow.examples.tutorials.mnist import input_data
  import torch
  
  class Net(torch.nn.Module):
      def __init__(self):
        super(Net, self).__init__()
        self.fc1 =  torch.nn.Linear(784, 100)
        self.fc2 =  torch.nn.Linear(100, 10)
  
      def forward(self, x):
        # x -&gt; (batch_size, 784)
        x = torch.relu(x)
        # x -&gt; (batch_size, 10)
        x = torch.softmax(x, dim=1)
        return x
  
  net = Net()
  net.zero_grad()
  Loss = torch.nn.CrossEntropyLoss()
  optimizer =  torch.optim.SGD(net.parameters(), lr=0.01)
  
  for epoch in range(1000):  # loop over the dataset multiple times
  
      batch_xs, batch_ys = mnist_m.train.next_batch(100)
      # convert to appropriate settins
      # note the input to the linear layer should be (n_sample, n_features)
      batch_xs = torch.tensor(batch_xs, requires_grad=True)
      # batch_ys -&gt; (batch_size,)
      batch_ys = torch.tensor(batch_ys, dtype=torch.int64)
  
      # forward
      # output -&gt; (batch_size, 10)
      output = net(batch_xs)
      # result -&gt; (batch_size,)
      result = torch.argmax(output, dim=1)
      loss = Loss(output, batch_ys)
  
      # backward
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
  </code></pre>