Empty val batch CUDA device fix (#7539)

Verified fix for https://github.com/ultralytics/yolov5/pull/7525#issuecomment-1106081123
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-04-22 12:01:14 -07:00
1 parent b804b36bc4
commit 813eba85b2
1 file changed
+1 -1
+1 -1
View File
@@ -227,7 +227,7 @@ def run(
if npr == 0:
if nl:
stats.append((correct, *torch.zeros((3, 0))))
stats.append((correct, *torch.zeros((3, 0), device=device)))
continue
# Predictions