Skip to content

Instantly share code, notes, and snippets.

@phildow
Last active January 21, 2020 20:03
Show Gist options
  • Save phildow/c43afc51438ec7390412e6a729e6f502 to your computer and use it in GitHub Desktop.
Save phildow/c43afc51438ec7390412e6a729e6f502 to your computer and use it in GitHub Desktop.
Prevent CUDA initialization failures in tensorflow with gpu_options.allow_growth=True
# Using tf.estimator without keras
session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth=True
run_config = tf.estimator.RunConfig(session_config=session_config)
classifier = tf.estimator.Estimator(model_fn=..., config=run_config)
# Using any keras
session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth=True
session = tf.Session(config=session_config)
tf.keras.backend.set_session(session)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment