autocast enable=torch.cuda.is_available() (#2748)
Browse files- models/common.py +1 -1
models/common.py
CHANGED
@@ -237,7 +237,7 @@ class autoShape(nn.Module):
|
|
237 |
return self
|
238 |
|
239 |
@torch.no_grad()
|
240 |
-
@torch.cuda.amp.autocast()
|
241 |
def forward(self, imgs, size=640, augment=False, profile=False):
|
242 |
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
243 |
# filename: imgs = 'data/samples/zidane.jpg'
|
|
|
237 |
return self
|
238 |
|
239 |
@torch.no_grad()
|
240 |
+
@torch.cuda.amp.autocast(torch.cuda.is_available())
|
241 |
def forward(self, imgs, size=640, augment=False, profile=False):
|
242 |
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
243 |
# filename: imgs = 'data/samples/zidane.jpg'
|