diff --git a/README.md b/README.md index a30075a2..193c834c 100644 --- a/README.md +++ b/README.md @@ -87,7 +87,7 @@ curl -fsSL https://raw.githubusercontent.com/containers/ramalama/s/install.sh | | Apple Silicon GPU (Linux / Asahi) | :white_check_mark: | | Apple Silicon GPU (macOS) | :white_check_mark: | | Apple Silicon GPU (podman-machine) | :white_check_mark: | -| Nvidia GPU (cuda) | :x: [Containerfile](https://github.com/containers/ramalama/blob/main/container-images/cuda/Containerfile) available but not published to quay.io | +| Nvidia GPU (cuda) | :white_check_mark: | | AMD GPU (rocm) | :white_check_mark: | ## COMMANDS diff --git a/ramalama/model.py b/ramalama/model.py index 086fde27..abe83e5a 100644 --- a/ramalama/model.py +++ b/ramalama/model.py @@ -101,8 +101,9 @@ def _image(self, args): gpu_type, _ = get_gpu() if gpu_type == "HIP_VISIBLE_DEVICES": return "quay.io/ramalama/rocm:latest" - - if gpu_type == "ASAHI_VISIBLE_DEVICES": + elif gpu_type == "CUDA_VISIBLE_DEVICES": + return "quay.io/ramalama/cuda:latest" + elif gpu_type == "ASAHI_VISIBLE_DEVICES": return "quay.io/ramalama/asahi:latest" return args.image