Skip to content

Commit

Permalink
autocast enable=torch.cuda.is_available() (ultralytics#2748)
Browse files Browse the repository at this point in the history
(cherry picked from commit fca5e2a)
  • Loading branch information
glenn-jocher authored and Lechtr committed Apr 12, 2021
1 parent d89f6f7 commit bfe2fb8
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down

0 comments on commit bfe2fb8

Please sign in to comment.