qurator-spk/eynollah
1
0
Fork 0
mirror of https://github.com/qurator-spk/eynollah.git synced 2026-05-26 07:39:22 +02:00
Code Issues Projects Releases Wiki Activity

model_zoo: also parse comma/colon syntax for device in Torch case

Browse source
This commit is contained in:
Robert Sachunsky 2026-05-12 20:02:40 +02:00
parent be4fe8c263
commit 17b311441a
1 changed files with 7 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -154,7 +154,7 @@ class EynollahModelZoo:
try:
gpus = tf.config.list_physical_devices('GPU')
if device:
if ',' in device:
if ':' in device:
for spec in device.split(','):
cat, dev = spec.split(':')
if fnmatchcase(model_category, cat):
@ -235,6 +235,12 @@ class EynollahModelZoo:
dev = torch.device('cpu')
if not device and torch.cuda.is_available():
device = 'GPU' # try
if device and ':' in device:
for spec in device.split(','):
cat, dev = spec.split(':')
if fnmatchcase('ocr', cat):
device = dev
break
if device and device.startswith('GPU'):
try:
dev = torch.device('cuda', int(device[3:] or 0))