Update to resolve FutureWarning from torch.load

[codecov[bot]]

Codecov Report

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

Project coverage is 99.60%. Comparing base (40987b5) to head (c3a1a7f). Report is 1 commits behind head on main.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #142 +/- ## ======================================= Coverage 99.60% 99.60% ======================================= Files 21 21 Lines 509 509 ======================================= Hits 507 507 Misses 2 2 ```

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