Pytorch beginner here. I am trying to predict the output of an AND gate using a neural network, but when I train the model, the loss doesn’t decrease and output stays at 0.2500 for every input. What am I doing wrong?
import torch 
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(2, 4)
    self.fc2 = nn.Linear(4,4)
    self.fc3 = nn.Linear(4, 1)
    
  def forward(self, x):
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = torch.sigmoid(self.fc3(x))
    
    return x

a = torch.tensor([[[0,0]],[[1,0]],[[0,1]],[[1,1]]], dtype = torch.float32)
y = torch.tensor([[0],[0],[0],[1]], dtype = torch.float32)

N = Net()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(N.parameters(), lr = 0.01)

for epoch in range(100):
  optimizer.zero_grad()
  out = N(a)
  loss = criterion(out, y)
  print(loss)
  loss.backward()
  optimizer.step()