FIX: track_gradients boolean bug.

sbi-dev · Apr 8, 2024 · ab67839 · ab67839
1 parent c40e4b1
commit ab67839
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sbi/samplers/mcmc/pymc_wrapper.py b/sbi/samplers/mcmc/pymc_wrapper.py
@@ -138,7 +138,7 @@ def __init__(
         self._device = device
 
         # create PyMC model object
-        track_gradients = step in (pymc.NUTS, pymc.HamiltonianMC)
+        track_gradients = step in ("nuts", "hmc")
         self._model = pymc.Model()
         potential = PyMCPotential(
             potential_fn, track_gradients=track_gradients, device=device