From cc99c36bbf40be0444d3c242c11af89b88aeb320 Mon Sep 17 00:00:00 2001
From: KaiyangZhou <k.zhou@qmul.ac.uk>
Date: Sat, 12 May 2018 17:45:16 +0100
Subject: [PATCH] add eval mode

---
 models/SqueezeNet.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/models/SqueezeNet.py b/models/SqueezeNet.py
index f302e56..6e472c0 100644
--- a/models/SqueezeNet.py
+++ b/models/SqueezeNet.py
@@ -108,6 +108,9 @@ class SqueezeNet(nn.Module):
         x10 = F.relu(self.conv10(x9))
         f = F.avg_pool2d(x10, x10.size()[2:]).view(x10.size(0), -1)
 
+        if not self.training:
+            return f
+
         if self.loss == {'xent'}:
             return f
         elif self.loss == {'xent', 'htri'}: