Enabling multiple GPUs

I am using a system containing two A4000 GPUs, however when I run the modulus examples only one of these is being used at a time. running the line:

torch.cuda.device_count()

in python within the docker container generates the output “2”, indicating both GPUs can be seen by pytorch (as far as I understand), but only GPU 0 is actually used during the running of the examples.

My question is what additional steps do I need to take to ensure that Modulus is able to run with both GPUs available on the system. The line I am using to start the container is:

docker run --gpus all --runtime nvidia --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it -p 8888:8888 modulus:22.09

Hi @bmschopp

You should be able to run Modulus in parallel using MPI and other typical parallel launching tools. For some information and example, have to look at the performance section of the docs or some of our examples that use multiple GPUs (e.g. Turbulent super resolution)