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:
1 parent
d40cd0d454
commit
61adf017f2
5 files changed
+8
-8
No files matched your search
@@ -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), [])
|
||||
|
||||
Reference in new issue
Block a user