Find free GPUs and set them as visible for pytorch or tensorflow

free_gpus.py

import os
import numpy as np

def get_free_gpus(n_gpus=2):
    os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
    memory_available = np.array([int(x.split()[2]) for x in open('tmp', 'r').readlines()])
    good_gpus = list((np.array(memory_available) > 6000).nonzero()[0])
    if len(good_gpus) == 0:
        raise IndexError
    good_gpus = good_gpus[:n_gpus]
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, good_gpus))