in models.py [0:0]
def forward(self, x):
for _, conv in enumerate(self.conv_layers):
out = conv(x)
x = F.max_pool2d(F.relu(out), self.pool_size, self.pool_size)
out = x.view(x.size(0), -1)
if self.normalize:
out = F.normalize(out)
return out