From b22f21cb6b68993c065d7540adc36fc96fbf82e7 Mon Sep 17 00:00:00 2001 From: MalteEbner Date: Mon, 4 Jan 2021 17:19:56 +0100 Subject: [PATCH] Bugfix in MultiBoxLoss when having no cuda gpu available --- layers/modules/multibox_loss.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/layers/modules/multibox_loss.py b/layers/modules/multibox_loss.py index fb49cf439..40436de12 100644 --- a/layers/modules/multibox_loss.py +++ b/layers/modules/multibox_loss.py @@ -32,8 +32,10 @@ class MultiBoxLoss(nn.Module): def __init__(self, num_classes, overlap_thresh, prior_for_matching, bkg_label, neg_mining, neg_pos, neg_overlap, encode_target, - use_gpu=True): + use_gpu=None): super(MultiBoxLoss, self).__init__() + if use_gpu is None: + use_gpu = torch.cuda.is_available() self.use_gpu = use_gpu self.num_classes = num_classes self.threshold = overlap_thresh