diff options
Diffstat (limited to 'gantools/biggan.py')
| -rw-r--r-- | gantools/biggan.py | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/gantools/biggan.py b/gantools/biggan.py index 571387b..71228ae 100644 --- a/gantools/biggan.py +++ b/gantools/biggan.py @@ -4,8 +4,20 @@ import tensorflow_hub as hub import numpy as np from itertools import cycle +#----------------------------------------------------------------- +# fix "could not create cudnn handle" error +# see: https://github.com/tensorflow/tensorflow/issues/24496 +from tensorflow.compat.v1 import ConfigProto +from tensorflow.compat.v1 import InteractiveSession +config = ConfigProto() +config.gpu_options.allow_growth = True +#----------------------------------------------------------------- + +session = InteractiveSession(config=config) + MODULE_PATH = 'https://tfhub.dev/deepmind/biggan-512/2' + class BigGAN(object): def __init__(self, module_path=MODULE_PATH): tf.reset_default_graph() |