Unverified 提交 fca5e2a4 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

autocast enable=torch.cuda.is_available() (#2748)

上级 c03d5903
......@@ -237,7 +237,7 @@ class autoShape(nn.Module):
return self
@torch.no_grad()
@torch.cuda.amp.autocast()
@torch.cuda.amp.autocast(torch.cuda.is_available())
def forward(self, imgs, size=640, augment=False, profile=False):
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
# filename: imgs = 'data/samples/zidane.jpg'
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论