replacing torch.cuda.set_device with

oandreeva-nv · oandreeva-nv · commit e1843571a29b · 2025-03-21T15:33:35.000-07:00
diff --git a/src/model.py b/src/model.py
@@ -320,7 +320,7 @@ def _validate_device_config(self):
                 f"Detected KIND_GPU model instance, explicitly setting GPU device={triton_device_id} for {triton_instance}"
             )
             # vLLM doesn't currently (v0.4.2) expose device selection in the APIs
-            torch.cuda.set_device(triton_device_id)
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(triton_device_id)
 
     def _setup_lora(self):
         self.enable_lora = False

Original file line number	Diff line number	Diff line change
`@@ -320,7 +320,7 @@ def _validate_device_config(self):`
`320`	`320`	`f"Detected KIND_GPU model instance, explicitly setting GPU device={triton_device_id} for {triton_instance}"`
`321`	`321`	`)`
`322`	`322`	`# vLLM doesn't currently (v0.4.2) expose device selection in the APIs`
`323`		`- torch.cuda.set_device(triton_device_id)`
	`323`	`+ os.environ["CUDA_VISIBLE_DEVICES"] = str(triton_device_id)`
`324`	`324`
`325`	`325`	`def _setup_lora(self):`
`326`	`326`	`self.enable_lora = False`