model.py

import torch
import torch.nn as nn

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        
        # Initialize the parent nn.Module class
        super(NeuralNet, self).__init__()

        # First hidden layer: input_size → hidden_size
        self.l1 = nn.Linear(input_size, hidden_size)

        # Second hidden layer: hidden_size → hidden_size
        self.l2 = nn.Linear(hidden_size, hidden_size)

        # Output layer: hidden_size → num_classes
        self.l3 = nn.Linear(hidden_size, num_classes)

        # Activation function (ReLU = Rectified Linear Unit)
        self.relu = nn.ReLU()

    def forward(self, x):
        # Pass input through first layer and apply ReLU
        out = self.l1(x)
        out = self.relu(out)

        # Pass through second layer and apply ReLU
        out = self.l2(out)
        out = self.relu(out)

        # Pass through output layer (no activation here)
        out = self.l3(out)

        # Note: No activation or softmax at the end
        # because we will apply softmax(CrossEntropyLoss) during training,
        # which expects raw logits.
        return out