togswr/docker-pytorch-jupyterlab

Using GPU in container.

Opened this issue · 0 comments

Hi, I found gpu not used after docker-compose up.

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

and the output is cpu.
After adding runtime: nvidia in docker-compose.yml, output changed to cuda.
i am not sure does this method to test whether gpu is used in container is right. Could you give some suggestion?