Using GPU in container.
Opened this issue · 0 comments
Li-Jun-SD commented
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?