I’ve been trying to plot the decision boundary of my neural network which I used for binary classification with the sigmoid function in the output layer but gor one error after another, I found many posts discussing the plotting of the decision boundary of a scikit-learn classifier but not a neural network built in PyTorch. Below is my neural network:
class NeuralNetwork(torch.nn.Module):
  def __init__(self):
  super(NeuralNetwork, self).__init__()
  self.fc1 = torch.nn.Linear(23, 16)
  self.fc2 = torch.nn.Linear(16, 14)
  self.fc3 = torch.nn.Linear(14, 10)
  self.fc4 = torch.nn.Linear(10, 5)
  self.fc5 = torch.nn.Linear(5, 1)

def forward(self, x):
  x = torch.relu(self.fc1(x))
  x = torch.relu(self.fc2(x))
  x = torch.relu(self.fc3(x))
  x = torch.relu(self.fc4(x))
  x = torch.sigmoid(self.fc5(x))
  return x

model = NeuralNetwork().double()

CUDA = torch.cuda.is_available()
if CUDA:
 model.cuda()

criterion = torch.nn.BCELoss(reduction='mean')
optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

model_1.train()

Precision = []
Cost = []

for epoch in range(10001):

  if CUDA:
    inputs = X_train.cuda()
    label = Y_train.cuda()
else:
  inputs = X_train
  label = Y_train

prediction = model_1(inputs)
loss = criterion(prediction, label)
accuracy = ((prediction > 0.5) == label).float().mean().item()

Cost.append(loss.item())
Precision.append(accuracy)

if epoch % 1000 == 0 or epoch == 30000:
  print("Epoch:", epoch, ",", "Loss:", loss.item(), ",", "Accuracy:", accuracy)

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

model_1.eval()

X_test = torch.from_numpy(X[27000:,:])
Y_test = torch.from_numpy(y[27000:,:]).double()

with torch.no_grad():

 y_pred = model_1(X_test)
print("Accuracy: ", ((y_pred > 0.5) == Y_test).float().mean().item())

As you can see my neural network has 23 features and I plan to choose only two features in my plot: X_test[:, 0] and X_test[:, 5].
I would greatly appreciate your help, thanks in advance.