From 38e3d33a6205cfb955887c2a3e30827c72fcd48d Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Mon, 1 Apr 2024 08:23:44 +0000
Subject: [PATCH] Add TPU to device config

---
 vllm/config.py | 3 +++
 1 file changed, 3 insertions(+)

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)