torch.empty() for speed improvements (#9025)

`torch.empty()` for speed improvement

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-08-18 20:12:33 +02:00
1 parent d40cd0d454
commit 61adf017f2
5 files changed
+8 -8

No files matched your search

+1 -1
View File
@@ -300,7 +300,7 @@ def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
try:
p = next(model.parameters()) # for device, type
imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz # expand
im = torch.zeros((1, 3, *imgsz)).to(p.device).type_as(p) # input image
im = torch.empty((1, 3, *imgsz)).to(p.device).type_as(p) # input image
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress jit trace warning
tb.add_graph(torch.jit.trace(de_parallel(model), im, strict=False), [])