From 94f59b51f8387aff83a205a8ba96b4a6093c5270 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 7 Oct 2021 20:35:13 -0700 Subject: [PATCH] Update --- utils/autobatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/autobatch.py b/utils/autobatch.py index 73ffe68b628f..3f97140c1a29 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -36,7 +36,7 @@ def autobatch(model, imgsz=640, fraction=0.9): f = t - (r + a) # free inside reserved print(f'{prefix}{t:.3g}G total, {r:.3g}G reserved, {a:.3g}G allocated, {f:.3g}G free') - batch_sizes = [1, 2, 4, 8, 16, 32] + batch_sizes = [1, 2, 4, 8, 16] try: img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes] y = profile(img, model, n=3, device=device)