text: import tensorflow as tf
  
  
  def limit_gpu(gpu_idx=0, mem=2 * 1024):
      gpus = tf.config.experimental.list_physical_devices('GPU')
      if gpus:
          # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
          try:
              # Use a single gpu
              tf.config.experimental.set_visible_devices(gpus[gpu_idx], 'GPU')
  
              # Limit memory
              tf.config.experimental.set_virtual_device_configuration(
                  gpus[0],
                  [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=mem)])  # 2 GB
              logical_gpus = tf.config.experimental.list_logical_devices('GPU')
              print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
          except RuntimeError as e:
              # Virtual devices must be set before GPUs have been initialized
              print(e)
  