From 98a3df0f8dd1f714107ae560bf1e66db173b3ccf Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Fri, 26 Apr 2024 08:56:26 +0000
Subject: [PATCH] Disable memory tracking

---
 vllm/worker/tpu_worker.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/vllm/worker/tpu_worker.py b/vllm/worker/tpu_worker.py
index 4e58830865758..d70224877ba8d 100644
--- a/vllm/worker/tpu_worker.py
+++ b/vllm/worker/tpu_worker.py
@@ -67,8 +67,8 @@ class TPUWorker(LoraNotSupportedWorkerBase):
                           os.path.expanduser("~/.vllm/jax_cache"))
 
         # DELETE
-        from jax_smi import initialise_tracking
-        initialise_tracking()
+        # from jax_smi import initialise_tracking
+        # initialise_tracking()
 
     def load_model(self):
         self.model_runner.load_model()