Skip to content

Commit 3331794

Browse files
gverkesakrishna1995
authored andcommitted
fix: make sure gpus are found in local_gpu run
1 parent c8b0876 commit 3331794

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

src/sagemaker/local/image.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -858,7 +858,13 @@ def _create_docker_host(
858858
# to setting --runtime=nvidia in the docker commandline.
859859
if self.instance_type == "local_gpu":
860860
host_config["deploy"] = {
861-
"resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
861+
"resources": {
862+
"reservations": {
863+
"devices": [
864+
{"count": "all", "capabilities": ["gpu"]}
865+
]
866+
}
867+
}
862868
}
863869

864870
if not self.is_studio and command == "serve":

0 commit comments

Comments
 (0)