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.
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.
Note that I intentionally use the dataset provided by TensorFlow to keep the data I use through 3 different methods consistent for accurate comparison.
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 -> (batch_size, 784)
      x = torch.relu(x)
      # x -> (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 -> (batch_size,)
    batch_ys = torch.tensor(batch_ys, dtype=torch.int64)

    # forward
    # output -> (batch_size, 10)
    output = net(batch_xs)
    # result -> (batch_size,)
    result = torch.argmax(output, dim=1)
    loss = Loss(output, batch_ys)

    # backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()