import tensorflow as tfgpus = tf.config.experimental.list_physical_devices(‘GPU’)if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e)

Source : How to limit GPU Memory in TensorFlow 2.0 (and 1.x) | by Jun-young, Cha | Medium