import torch
import torch.nn as nn


class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.l1 = nn.Linear(input_size, hidden_size) 
        self.l2 = nn.Linear(hidden_size, hidden_size) 
        self.l3 = nn.Linear(hidden_size, hidden_size)
        self.l4 = nn.Linear(hidden_size, num_classes)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        out = self.l1(x)
        out = self.relu(out)
        out = self.l2(out)
        out = self.relu(out)
        out = self.l3(out)
        out = self.relu(out)
        out = self.l4(out)
        # no activation and no softmax at the end
        return out
[(8, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (9, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (10, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (11, 'torch.nn.Linear', 'nn.Linear', True, 'import torch.nn as nn\n'), (12, 'torch.nn.ReLU', 'nn.ReLU', True, 'import torch.nn as nn\n')]