PyTorch implementation of 2 linear layers

Demo1.py

# PyTorch implementation of 2 linear layers

def run():
    ...
    train_loader = torch.utils.data.DataLoader(...)

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.fc1 = nn.Linear(784, 50)
            self.fc2 = nn.Linear(50, 10)

        def forward(self, x):
            x = x.view(-1, 784)
            x = F.relu(self.fc1(x))
            x = self.fc2(x)
            return F.log_softmax(x, dim=1)

    model = Net()
    optimizer = optim.SGD(...)

    def train(epoch):
        for batch_idx, (data, target) in enumerate(train_loader):
            ...
            loss.backward()
            optimizer.step()
            if (((batch_idx + 1) % args.log_interval) == 0):
                print_time_and_loss()

    for epoch in range(1, args.epochs + 1):
        train(epoch)