116 make plot and predict consistent

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 96.65%. Comparing base (f7bf0f4) to head (e88f54a). Report is 7 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #118 +/- ## ========================================== + Coverage 96.64% 96.65% +0.01% ========================================== Files 22 22 Lines 655 658 +3 ========================================== + Hits 633 636 +3 Misses 22 22 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.