Guard CUDA calls with an explicit check (#516)

RangiLyu · Jul 12, 2023 · 3c9607c · 3c9607c
1 parent d3fb34f
commit 3c9607c
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/nanodet/model/arch/one_stage_detector.py b/nanodet/model/arch/one_stage_detector.py
@@ -47,14 +47,23 @@ def forward(self, x):
 
  def inference(self, meta):
  with torch.no_grad():
- torch.cuda.synchronize()
+ is_cuda_available = torch.cuda.is_available()
+ if is_cuda_available:
+ torch.cuda.synchronize()
+
  time1 = time.time()
  preds = self(meta["img"])
- torch.cuda.synchronize()
+
+ if is_cuda_available:
+ torch.cuda.synchronize()
+
  time2 = time.time()
  print("forward time: {:.3f}s".format((time2 - time1)), end=" | ")
  results = self.head.post_process(preds, meta)
- torch.cuda.synchronize()
+
+ if is_cuda_available:
+ torch.cuda.synchronize()
+
  print("decode time: {:.3f}s".format((time.time() - time2)), end=" | ")
  return results