glenn-jocher commited on
Commit
fca5e2a
·
unverified ·
1 Parent(s): c03d590

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

Browse files
Files changed (1) hide show
  1. 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'