Merge pull request #525 from ztqakita/master

FIx bug in Delay
brainpy · Oct 30, 2023 · d085391 · d085391
2 parents e6c6664 + d1a637d
commit d085391
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/brainpy/_src/delay.py b/brainpy/_src/delay.py
@@ -389,7 +389,11 @@ def _init_data(self, length: int, batch_size: int = None):
     else:
       batch_axis = self.target.batch_axis + 1
 
-    f = jax.jit(jnp.zeros, static_argnums=0, static_argnames='dtype', out_shardings=self.sharding)
+    if self.sharding is None:
+      f = jnp.zeros
+    else:
+      f = jax.jit(jnp.zeros, static_argnums=0, static_argnames='dtype', out_shardings=self.sharding)
+
     data = f((length,) + self.target.shape, dtype=self.target.dtype)
     if self.data is None:
       self.data = bm.Variable(data, batch_axis=batch_axis)