Fix sum() bug in multibox_loss that was due to old pytorch compatibility

Author: amdegroot

layers/modules/multibox_loss.py

@@ -6,6 +6,7 @@
 from data import v2 as cfg
 from ..box_utils import match, log_sum_exp
 
+
 class MultiBoxLoss(nn.Module):
     """SSD Weighted Loss Function
     Compute Targets:
@@ -79,7 +80,7 @@ def forward(self, predictions, targets):
         conf_t = Variable(conf_t, requires_grad=False)
 
         pos = conf_t > 0
-        num_pos = pos.sum(keepdim=True)
+        num_pos = pos.sum(dim=1, keepdim=True)
 
         # Localization Loss (Smooth L1)
         # Shape: [batch,num_priors,4]

ssd.py

@@ -112,7 +112,8 @@ def load_weights(self, base_file):
         other, ext = os.path.splitext(base_file)
         if ext == '.pkl' or '.pth':
             print('Loading weights into state dict...')
-            self.load_state_dict(torch.load(base_file, map_location=lambda storage, loc: storage))
+            self.load_state_dict(torch.load(base_file,
+                                 map_location=lambda storage, loc: storage))
             print('Finished!')
         else:
             print('Sorry only .pth and .pkl files supported.')
@@ -199,7 +200,7 @@ def build_ssd(phase, size=300, num_classes=21):
     if size != 300:
         print("Error: Sorry only SSD300 is supported currently!")
         return
-    base_,extras_,head_=multibox(vgg(base[str(size)], 3),
-                                 add_extras(extras[str(size)], 1024),
-                                 mbox[str(size)], num_classes)   
-    return SSD(phase,base_,extras_,head_, num_classes)
+    base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
+                                     add_extras(extras[str(size)], 1024),
+                                     mbox[str(size)], num_classes)
+    return SSD(phase, base_, extras_, head_, num_classes)