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

This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-04-09 13:34:49 +02:00
1 parent c03d590320
commit fca5e2a48f
1 file changed
+1 -1
+1 -1
View File
@@ -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'