diff --git a/torchx/schedulers/docker_scheduler.py b/torchx/schedulers/docker_scheduler.py index b56e4381b..3dac9f81a 100644 --- a/torchx/schedulers/docker_scheduler.py +++ b/torchx/schedulers/docker_scheduler.py @@ -314,7 +314,7 @@ def _submit_dryrun(self, app: AppDef, cfg: DockerOpts) -> AppDryRunInfo[DockerJo c.kwargs["device_requests"] = [ DeviceRequest( count=resource.gpu, - capabilities=[["compute"]], + capabilities=[["compute", "utility"]], ) ] req.containers.append(c) diff --git a/torchx/schedulers/test/docker_scheduler_test.py b/torchx/schedulers/test/docker_scheduler_test.py index 846cd190a..14ab3ee25 100644 --- a/torchx/schedulers/test/docker_scheduler_test.py +++ b/torchx/schedulers/test/docker_scheduler_test.py @@ -91,7 +91,7 @@ def test_submit_dryrun(self) -> None: "device_requests": [ DeviceRequest( count=4, - capabilities=[["compute"]], + capabilities=[["compute", "utility"]], ) ], "devices": [