Is it possible to make JukeBox use more of my GPUs?

[ratsaby]

I installed JukeBox and it is running straight from the Jupiter notebook.

I would like to know if it is possible to make the code use more than one GPU?

I am unfamiliar with PyTorch, but I could see in this regards, the following which seems to cause

JukeBox to run on a single GPU (whose ID is 0)

==================================
def setup_dist_from_mpi(
master_addr="127.0.0.1", backend="nccl", port=29500, n_attempts=5, verbose=False
):
if dist.is_available():
return _setup_dist_from_mpi(master_addr, backend, port, n_attempts, verbose)
else:
use_cuda = torch.cuda.is_available()
print(f'Using cuda {use_cuda}')

    mpi_rank = 0
    local_rank = 0

    device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
    torch.cuda.set_device(local_rank)

    return mpi_rank, local_rank, device

=====================================

Can you advise how can I make it use all my 4 GPUs ?

Thanks
J

[prafullasd]

The usual way of running the sample.py script with multiple GPU's would be to launch with mpiexec -n {ngpus} python sample.py ...
(also make sure to change logdir here to {hps.name}_{dist.get_rank()}/level_{level} so the ranks dont overwrite each other)

For local jupyter notebooks, you can try running 4 separate notebooks, and for each notebook launch it with CUDA_VISIBLE_DEVICES={gpu_id} jupyter notebook so that it uses that specific GPU.