Просмотр исходного кода

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

5.0
Glenn Jocher GitHub 3 лет назад
Родитель
Сommit
fca5e2a48f
Не найден GPG ключ соответствующий данной подписи Идентификатор GPG ключа: 4AEE18F83AFDEB23
1 измененных файлов: 1 добавлений и 1 удалений
  1. +1
    -1
      models/common.py

+ 1
- 1
models/common.py Просмотреть файл

@@ -237,7 +237,7 @@ class autoShape(nn.Module):
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'

Загрузка…
Отмена
Сохранить