qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-06-16 09:59:13 +02:00
Code Issues Projects Releases Wiki Activity

ModelZoo device selection: warn if model category still unmatched…

Browse source 
(and try GPU)
This commit is contained in:
Robert Sachunsky 2026-06-12 14:58:32 +02:00
parent 94082bc64a
commit 60c9f4786c
1 changed files with 11 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

13
src/eynollah/model_zoo/model_zoo.py
View file

@ -182,7 +182,10 @@ class EynollahModelZoo:
if fnmatchcase(model_category, cat): if fnmatchcase(model_category, cat):
device = dev device = dev
break break
if device == 'CPU': if ':' in device:
self.logger.warning("missing device specification for model type %s", model_category)
gpus = gpus[:1]
elif device == 'CPU':
gpus = [] gpus = []
else: else:
assert device.startswith('GPU') assert device.startswith('GPU')
@ -227,6 +230,9 @@ class EynollahModelZoo:
if fnmatchcase('ocr', cat): if fnmatchcase('ocr', cat):
device = dev device = dev
break break
if ':' in device:
self.logger.warning("missing device specification for model type %s", model_category)
device = 'GPU'
if device and device.startswith('GPU'): if device and device.startswith('GPU'):
try: try:
device0 = torch.device('cuda', int(device[3:] or 0)) device0 = torch.device('cuda', int(device[3:] or 0))
@ -309,7 +315,10 @@ class EynollahModelZoo:
if fnmatchcase(model_category, cat): if fnmatchcase(model_category, cat):
device = dev device = dev
break break
if device == 'CPU': if ':' in device:
self.logger.warning("missing device specification for model type %s", model_category)
gpu = 0
elif device == 'CPU':
gpu = -1 gpu = -1
else: else:
assert device.startswith('GPU') assert device.startswith('GPU')