enable inference on cpu (facebookresearch#821)

now can inference on cpu by setting MODEL.DEVICE cpu
Lyears · May 25, 2019 · a5f77c0 · a5f77c0
1 parent 0b2a075
commit a5f77c0
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(