Update usages of nn.Softmax and torch.mul_ to avoid deprecation warnings

amdegroot · Feb 15, 2018 · d15c97f · d15c97f
1 parent afebb30
commit d15c97f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 31 deletions.
diff --git a/demo/demo.ipynb b/demo/demo.ipynb
diff --git a/ssd.py b/ssd.py
@@ -43,7 +43,7 @@ def __init__(self, phase, base, extras, head, num_classes):
         self.conf = nn.ModuleList(head[1])
 
         if phase == 'test':
-            self.softmax = nn.Softmax()
+            self.softmax = nn.Softmax(dim=-1)
             self.detect = Detect(num_classes, 0, 200, 0.01, 0.45)
 
     def forward(self, x):
@@ -97,8 +97,8 @@ def forward(self, x):
         if self.phase == "test":
             output = self.detect(
                 loc.view(loc.size(0), -1, 4),                   # loc preds
-                self.softmax(conf.view(-1, self.num_classes)) \
-                    .view(conf.size(0), -1, self.num_classes),  # conf preds
+                self.softmax(conf.view(conf.size(0), -1,
+                    self.num_classes)),                         # conf preds
                 self.priors.type(type(x.data))                  # default boxes
             )
         else: