Fixes predictions with structured probabilistic model

[ziatdinovmax]

Addresses breaking change in https://github.com/pyro-ppl/numpyro/releases/tag/0.14.0
(Note: BNNs were also affected by it, but now all should work fine)

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 96.55172% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 95.29%. Comparing base (14d1b50) to head (784a6f3).

Files	Patch %	Lines
gpax/models/spm.py	93.75%	1 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #93   +/-   ##
=======================================
  Coverage   95.29%   95.29%           
=======================================
  Files          55       55           
  Lines        4247     4269   +22     
=======================================
+ Hits         4047     4068   +21     
- Misses        200      201    +1     

Flag	Coverage Δ
unittests	95.29% <96.55%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

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