Adapt fit_hmc to work with latest blackjax version#14

RaulPL wants to merge 1 commit into

probml:mainprobml/sts-jax:mainfrom

RaulPL:fix_samplingRaulPL/sts-jax:fix_samplingCopy head branch name to clipboard