diff options
| author | Vee9ahd1 <[email protected]> | 2019-09-22 19:43:03 -0400 |
|---|---|---|
| committer | Christian <[email protected]> | 2019-09-22 19:46:11 -0400 |
| commit | 044232a0947fcf02118a268e45a615cc1a8065a2 (patch) | |
| tree | 213e6375107eee8cb5cf66dc4fbdc661200f0cb1 /gantools/biggan.py | |
| parent | ed90828eb9712de21ecf7bfa5dfe88bd27672a0b (diff) | |
patched for artbreeder support and fixed some weird CUDNN_STATUS_INTERNAL_ERROR. if CPU support is broken now, this commit is probably why.
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() |