diff --git a/vllm/config.py b/vllm/config.py index a901729194c9b..8dd670a4eccf8 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -596,6 +596,9 @@ class DeviceConfig: # Some device types require processing inputs on CPU if self.device_type in ["neuron"]: self.device = torch.device("cpu") + elif self.device_type in ["tpu"]: + # Will be set by `xm.xla_device()` + self.device = None else: # Set device with device type self.device = torch.device(self.device_type)