fix jnp usage

AI-Hypercomputer · Dec 3, 2024 · ee48797 · ee48797
1 parent 35e88f2
commit ee48797
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pathwaysutils/persistence/helper.py b/pathwaysutils/persistence/helper.py
@@ -172,7 +172,7 @@ def get_read_request(
 def get_bulk_read_request(
     location_path: str,
     names: str,
-    dtypes: jnp.dtype,
+    dtypes: np.dtype,
     shapes: Sequence[Sequence[int]],
     shardings: Sequence[jax.sharding.Sharding],
     devices: Sequence[jax.Device],