enable inference on cpu (#821)

now can inference on cpu by setting MODEL.DEVICE cpu
facebookresearch · May 25, 2019 · d7e3c65 · d7e3c65
1 parent cd35ff0
commit d7e3c65
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/maskrcnn_benchmark/engine/inference.py b/maskrcnn_benchmark/engine/inference.py
@@ -29,7 +29,8 @@ def compute_on_dataset(model, data_loader, device, timer=None):
             else:
                 output = model(images.to(device))
             if timer:
-                torch.cuda.synchronize()
+                if not cfg.MODEL.DEVICE == 'cpu':
+                    torch.cuda.synchronize()
                 timer.toc()
             output = [o.to(cpu_device) for o in output]
         results_dict.update(