geekswaroop
diff --git a/‎eval.py
+7-2 b/‎eval.py
+7-2
diff --git a/‎static/input.png
-10.8 KB b/‎static/input.png
-10.8 KB
diff --git a/‎static/output_densenet.png
-1.92 KB b/‎static/output_densenet.png
-1.92 KB
diff --git a/‎static/output_resnet101.png
-5.02 KB b/‎static/output_resnet101.png
-5.02 KB
diff --git a/‎static/output_resnet152.png
-4.44 KB b/‎static/output_resnet152.png
-4.44 KB
diff --git a/‎static/output_resnet18.png
-2.95 KB b/‎static/output_resnet18.png
-2.95 KB
diff --git a/‎static/output_resnet34.png
-2.68 KB b/‎static/output_resnet34.png
-2.68 KB
diff --git a/‎static/output_resnet50.png
-2.16 KB b/‎static/output_resnet50.png
-2.16 KB
@@ -57,9 +57,14 @@ def build_network(snapshot, backend, gpu):
 
     if snapshot is not None:
         _, epoch = os.path.basename(snapshot).split('_')
-        if not epoch == 'last':
+        if epoch != 'last':
             epoch = int(epoch)
-        net.load_state_dict(torch.load(snapshot))
+
+        if gpu:
+            net.load_state_dict(torch.load(snapshot))
+        else:
+            net.load_state_dict(torch.load(snapshot, map_location=torch.device('cpu')))
+
         logging.info("Snapshot for epoch {} loaded from {}".format(epoch, snapshot))
 
     if gpu: