qurator-spk
/
eynollah
mirror of https://github.com/qurator-spk/eynollah.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

set_memory_growth to all GPU devices alike

Browse Source
pull/100/head
Robert Sachunsky 2 years ago committed by GitHub
parent fd56b86acf
commit 8fe3567123
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File

3
qurator/eynollah/eynollah.py
Unescape Escape View File

@ -496,7 +496,8 @@ class Eynollah:
#session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) #session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))
physical_devices = tf.config.list_physical_devices('GPU') physical_devices = tf.config.list_physical_devices('GPU')
try: try:
tf.config.experimental.set_memory_growth(physical_devices[0], True) for device in physical_devices:
tf.config.experimental.set_memory_growth(device, True)
except: except:
self.logger.warning("no GPU device available") self.logger.warning("no GPU device available")
if model_dir.endswith('.h5') and Path(model_dir[:-3]).exists(): if model_dir.endswith('.h5') and Path(model_dir[:-3]).exists():

Write Preview
Loading…
Cancel
Save