Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
models/common.py
+1 -1
View File
@@ -237,7 +237,7 @@ class autoShape(nn.Module):
return self return self
@torch.no_grad() @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): def forward(self, imgs, size=640, augment=False, profile=False):
# Inference from various sources. For height=640, width=1280, RGB images example inputs are: # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
# filename: imgs = 'data/samples/zidane.jpg' # filename: imgs = 'data/samples/zidane.jpg'