I’m not aware of a straightforward way, so please correct me, if there is a simpler solution.
For now you could call nvidia-smi inside your script and based on this output select the GPU you would like to use. The returned memory stats will be wrapped into a pandas.DataFrame so that you can sort it etc.
Here is a sample code:
import subprocess
import sys
if sys.version_info[0] < 3: 
    from StringIO import StringIO
else:
    from io import StringIO
import pandas as pd

def get_free_gpu():
    gpu_stats = subprocess.check_output(["nvidia-smi", "--format=csv", "--query-gpu=memory.used,memory.free"])
    gpu_df = pd.read_csv(StringIO(gpu_stats),
                         names=['memory.used', 'memory.free'],
                         skiprows=1)
    print('GPU usage:\n{}'.format(gpu_df))
    gpu_df['memory.free'] = gpu_df['memory.free'].map(lambda x: x.rstrip(' [MiB]'))
    idx = gpu_df['memory.free'].idxmax()
    print('Returning GPU{} with {} free MiB'.format(idx, gpu_df.iloc[idx]['memory.free']))
    return idx

free_gpu_id = get_free_gpu()
torch.cuda.set_device(free_gpu_id)