From ee90d08b229ec7cfa84235f7878a0768da1bf3f6 Mon Sep 17 00:00:00 2001
From: Fanhai Lu <154379058+FanhaiLu1@users.noreply.github.com>
Date: Wed, 24 Apr 2024 15:43:08 -0700
Subject: [PATCH] do nothing for nd array (#52)

---
 jetstream/engine/engine_api.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/jetstream/engine/engine_api.py b/jetstream/engine/engine_api.py
index d8289b9b..bbf64bb4 100644
--- a/jetstream/engine/engine_api.py
+++ b/jetstream/engine/engine_api.py
@@ -82,6 +82,9 @@ class ResultTokens(abc.ABC):
 
   def copy_to_host_async(self: "ResultTokens") -> None:
     """Copy to host asynchronously."""
+    # Do nothing for np array
+    if isinstance(self.data, np.ndarray):
+      return
     self.data.copy_to_host_async()
 
   def convert_to_numpy(self: "ResultTokens") -> "ResultTokens":