From 2aa9831dd380ccfcf219068ee015e00a018fa4ca Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Thu, 25 Apr 2024 23:40:44 +0000
Subject: [PATCH] Minor

---
 vllm/model_executor/models/jax/ops/write_to_cache.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/vllm/model_executor/models/jax/ops/write_to_cache.py b/vllm/model_executor/models/jax/ops/write_to_cache.py
index 8124bd6fd2..cdd562afcd 100644
--- a/vllm/model_executor/models/jax/ops/write_to_cache.py
+++ b/vllm/model_executor/models/jax/ops/write_to_cache.py
@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import chex
 import jax
 import jax.numpy as jnp