class LeNet5(nn.Module): def __init__(self): super(LeNet5, self).__init__() # define a 2D convolutional layer self.hidden_1 = torch.nn.Conv2d(1,6,kernel_size=5,stride=1,padding=2) # define a maxpool layer self.hidden_2 = nn.MaxPool2d(2, stride=2) # new 2D convolutional layer self.hidden_3 = torch.nn.Conv2d(6,16,kernel_size=5,stride=1) # another maxpool layer self.hidden_4 = nn.MaxPool2d(2, stride=2) # first linear layer self.hidden_5 = nn.Linear(16*5*5,120, bias=True) # second linear layer self.hidden_6 = nn.Linear(120,84, bias=True) # final output layer self.output = nn.Linear(84, 10, bias=False) # activation function self.activation = nn.ReLU() def forward(self, x): for i in range(1,7): if i==5: x=x.flatten(start_dim=1) # activate pass through the first layer f = getattr(self,f"hidden_{i}") x = self.activation(f(x)) return self.output(x) # return output x = torch.randn((1, 1, 28, 28)) model = LeNet5() y = model(x) print(y) print(model)