yolov3: reduce batch size due to OOM

pytorch · Oct 23, 2023 · fe396c3 · fe396c3
1 parent b5631ed
commit fe396c3
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/torchbenchmark/models/yolov3/__init__.py b/torchbenchmark/models/yolov3/__init__.py
@@ -30,7 +30,7 @@ class Model(BenchmarkModel):
     task = COMPUTER_VISION.SEGMENTATION
     # Original train batch size: 16
     # Source: https://github.com/ultralytics/yolov3/blob/master/train.py#L447
-    DEFAULT_TRAIN_BSIZE = 16
+    DEFAULT_TRAIN_BSIZE = 8 # reduced from 16 to 8 due to OOM in TorchInductor dashboard
     DEFAULT_EVAL_BSIZE = 8
     # yolov3 CUDA inference test uses amp precision
     DEFAULT_EVAL_CUDA_PRECISION = "amp"
@@ -102,6 +102,4 @@ def eval(self) -> Tuple[torch.Tensor]:
     def device_str(self):
         """YoloV3 uses individual GPU indices."""
         return str(
-            torch.cuda.current_device() if self.device == "cuda"
-            else self.device
-        )
+            torch.cuda.current_device() if