Fix pytorch issue

[HCookie]

Closes #577

• Add to .cpu()

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.12%. Comparing base (aec3c03) to head (8a53865).
Report is 8 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop     #578   +/-   ##
========================================
  Coverage    90.12%   90.12%           
========================================
  Files          151      151           
  Lines        11165    11165           
  Branches       529      529           
========================================
  Hits         10063    10063           
  Misses         933      933           
  Partials       169      169           

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