fix

netket · May 1, 2024 · 6e1716f · 6e1716f
1 parent 8fc9827
commit 6e1716f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/netket/sampler/rules/hamiltonian.py b/netket/sampler/rules/hamiltonian.py
@@ -105,7 +105,7 @@ def __transition(v, rand_vec):
  @partial(sharding_decorator, sharded_args_tree=(True, True))
  def _transition(v, rand_vec):
  return jax.pure_callback(
- _transition,
+ __transition,
  (
  jax.core.ShapedArray(v.shape, v.dtype),
  jax.core.ShapedArray((v.shape[0],), log_prob_dtype),