def forward()

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