Fix remaining issues

[pat-alt]

• gradients in mcmc_samples should be computed for batches, not the whole dataset.
• for this to work, it probably does not make sense to have inputs as part of model.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 94.85%. Comparing base (50c1fd6) to head (3f4ff73).
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #14      +/-   ##
==========================================
+ Coverage   94.35%   94.85%   +0.49%     
==========================================
  Files           5        5              
  Lines         124      136      +12     
==========================================
+ Hits          117      129      +12     
  Misses          7        7              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.