load weights to cpu if use_cuda is False

Author: namangup

efficientdet_test.py

@@ -58,7 +58,10 @@
 
 model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
                              ratios=anchor_ratios, scales=anchor_scales)
-model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
+if use_cuda:
+    model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth'))
+else:
+    model.load_state_dict(torch.load(f'weights/efficientdet-d{compound_coef}.pth', map_location=torch.device('cpu')))
 model.requires_grad_(False)
 model.eval()