Hi!
I am working on clothing detection and instead of k-nearest neighbours I decided to learn how neural networks work. So I created very simple network and started to train it based on a tutorial: https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html 5
At first it looked very primising. The loss was decreasing in a satisfying rate, but it stopped as quickly as it started.
As I am at beginner in neural networks I don’t know what to do to imporve my neural network now.
Here is the code:
import torchvision
import torchvision.transforms as transforms
import torch
from torch import nn
import torch.optim as optim
import pickle as pkl


if __name__ == '__main__':

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data = pkl.load(open('train.pkl', mode='rb'))
    traindata = data[0]
    labels = data[1]
    tensor_data = torch.stack([torch.Tensor(i) for i in traindata])
    tensor_labels = torch.from_numpy(labels)

    dataset = torch.utils.data.TensorDataset(tensor_data, tensor_labels)

    trainloader = torch.utils.data.DataLoader(dataset, batch_size=16,
                                              shuffle=True, num_workers=2)
    testloader = torch.utils.data.DataLoader(dataset, batch_size=16,
                                             shuffle=False, num_workers=2)

    torch.multiprocessing.freeze_support()


    class SimpleNetwork(nn.Module):
        def __init__(self):
            super(SimpleNetwork, self).__init__()
            self.fc1 = nn.Linear(1296, 10)
            self.sig = nn.Sigmoid()
            self.params = list(self.fc1.parameters())

        def forward(self, x):
            x = self.fc1(x)
            x = self.sig(x)
            return x

        def get_params(self):
            w = self.params
            return w


    net = SimpleNetwork()
    #net.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

    print('testing started')

    for epoch in range(20):
        net.load_state_dict(torch.load('parameters.pkl'))
        net.eval()
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            #inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            #print(torch.sum(net.fc1.weight.data))

            running_loss += loss.item()
            if i % 2000 == 1999:
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

    data = net.fc1.weight.data.numpy(), net.fc1.bias.data.numpy()
    with open('weights.pkl', 'wb') as f:
        pkl.dump(data, f)

    torch.save(net.state_dict(), 'parameters.pkl')
    print('finished testing')



My data (train.pkl) is a file containing a numpy array of images (represented as 1296 floats) and a array of labels (from 0 to 9)
I need to extract weights from this model, hence last few lines of code
What can I do make this neural network learn?
Please help