do nothing for nd array (#52)

AI-Hypercomputer · Apr 24, 2024 · ee90d08 · ee90d08
1 parent a0df320
commit ee90d08
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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":