import tensorflow as tf

    memory_limit = int(fraction*total_memory)
    print(memory_info)
    if tf.version.VERSION[0]=="2":
        gpus = tf.config.experimental.list_physical_devices('GPU')
        tf.config.experimental.set_memory_growth(gpus[0], True)
        tf.config.experimental.set_virtual_device_configuration(gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=memory_limit)])
    else:
        gpu_options = tf.GPUOptions(allow_growth=allow_growth,
                                  per_process_gpu_memory_fraction=fraction)
        config = tf.ConfigProto(gpu_options=gpu_options)
        session = tf.Session(config=config)
