mirror of
https://github.com/qurator-spk/eynollah.git
synced 2025-06-09 20:29:55 +02:00
Merge pull request #100 from bertsky/patch-2
set_memory_growth to all GPU devices alike
This commit is contained in:
commit
52d2e0b098
1 changed files with 2 additions and 1 deletions
|
@ -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():
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue