inference_on_dataset make sure the device is a gpu
Summary:
Before calling `torch.cuda.synchronize()`, need to make sure running a gpu.
Pull Request resolved:
@@ -119,6 +119,7 @@ def inference_on_dataset(model, data_loader, evaluator):
start_compute_time = time.time()
outputs = model(inputs)
+ if torch.cuda.is_available():
torch.cuda.synchronize()
total_compute_time += time.time() - start_compute_time
evaluator.process(inputs, outputs)
