Update ssd.py

3b0eb1f6 · JiaQi Xu · GitHub · ccb1ac0b · 3b0eb1f6
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

ssd.py ssd.py +1 -1

未找到文件。
--- a/ssd.py
+++ b/ssd.py
@@ -55,9 +55,9 @@ class SSD(object):
        model = ssd.get_ssd("test",self.num_classes)
        model.load_state_dict(torch.load(self.model_path))
        self.net = model.eval()
-        self.net = torch.nn.DataParallel(self.net)

        if self.cuda:
+            self.net = torch.nn.DataParallel(self.net)
            cudnn.benchmark = True
            self.net = self.net.cuda()