In trying to plot the the response curves I get an error

[PabloAguirreSolana]

In trying to plot the response curves I get this error, every single time I run it. And frankly don't know how to fix it.

TracerIntegerConversionError: The index() method was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=3/0)>

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError