From 5ef589361ae68f734401166a20044ec09784f081 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 9 Apr 2021 13:34:49 +0200 Subject: [PATCH] autocast enable=torch.cuda.is_available() (#2748) --- models/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/common.py b/models/common.py index 412e9bf1e411..c77ecbeceace 100644 --- a/models/common.py +++ b/models/common.py @@ -237,7 +237,7 @@ def autoshape(self): 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'