Hey, i have some training data in the format:
X = [[0.1, 0.2], [0.1, 0.3], …] - (x,y) input data
y = [[1,0,0], [0,1,0], … ] - one-hot encoded output
I have defined the model and training process as the following:
w0 = torch.tensor(np.random.randn(2,64), requires_grad=True)
w1 = torch.tensor(np.random.randn(64,3), requires_grad=True)

optim = torch.optim.Adam([w0, w1])

for epoch in range(1000):
        out = torch.tensor(X, requires_grad=True).matmul(w0).relu().matmul(w1)
        out = torch.nn.functional.log_softmax(out, dim=1)

        y_pred = out.detach().numpy().argmax(axis=1)
        acc = np.mean(y_pred == y_true) # y_true is the categorical representation

        # Categorical cross-entropy loss
        loss = -(y * out).sum(dim=1).mean()

        optim.zero_grad()
        loss.backward()
        optim.step()

        print(f"loss: {loss.data}, acc: {acc}")

The accuracy is stuck at 40% when it should be around 90%.