`class Net(nn.Module):
def init(self, num_channels):
super(Net,self).init()
self.num_channels = num_channels
self.conv1 = nn.Conv2d(3, self.num_channels, 3, stride = 1, padding = 1)
self.conv2 = nn.Conv2d(self.num_channels, self.num_channels*2, 3, stride = 1, padding = 1)
self.conv3 = nn.Conv2d(self.num_channels*2, self.num_channels*4, 3, stride = 1, padding = 1)
self.fc1 = nn.Linear(self.num_channels*4*8*8, self.num_channels*4)
self.fc2 = nn.Linear(self.num_channels*4, 6)
def forward(self,x):
#Empieza 3x64x64
x = self.conv1(x) # num_channels x 64 x 64
x = F.relu(F.max_pool2d(x, 2)) # num_channels x 32 x 32
x = self.conv2(x) # num_channels2 x 32 x32
x = F.relu(F.max_pool2d(x, 2)) #num_channels2 x 16 x 16
x = self.conv3(x) # num_channels4 x16x16
x = F.relu(F.max_pool2d(x, 2)) # num_channels4 x 8 x 8
#flatten
x = x.view(-1, self.num_channels*4*8*8)
#fc
x = self.fc1(x)
x = F.relu(x)
x = self.fc2(x)
#log_softmax
x = F.log_softmax(x, dim=1)
return x`
