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

(cherry picked from commit fca5e2a)
Lechtr · Apr 12, 2021 · bfe2fb8 · bfe2fb8
1 parent d89f6f7
commit bfe2fb8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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'