
import torch 
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
 #         self.conv1 = nn.Conv2d(1, 20, 5, 1)
#         self.conv2 = nn.Conv2d(20, 50, 5, 1)
#         self.fc1 = nn.Linear(4*4*50, 500)
#         self.fc2 = nn.Linear(500, 10) *)

        self.features = nn.Sequential(nn.Conv2d(1, 20, 5, 1), nn.ReLU(True), nn.MaxPool2d(2,2),
                                      nn.Conv2d(20, 50, 5, 1), nn.ReLU(True), nn.MaxPool2d(2,2))
        self.classifier = nn.Sequential(nn.Linear(4*4*50, 500), nn.ReLU(True),
                                        nn.Linear(500, 10))

    def forward(self, x):
#         x = F.relu(self.conv1(x))
#         x = F.max_pool2d(x, 2, 2)
#         x = F.relu(self.conv2(x))
#         x = F.max_pool2d(x, 2, 2)
#         x = x.view(-1, 4*4*50)
#         x = F.relu(self.fc1(x))
#         x = self.fc2(x) *)
        x = self.features(x)
        x = x.view(-1,4*4*50)
        x = self.classifier(x)
        return x#F.log_softmax(x, dim=1)
    

 