import tensorflow as tf physical_devices = tf.config.experimental.list_physical_devices('GPU') config = tf.config.experimental.set_memory_growth(physical_devices[0], True)