rm unnecessary device put

[miguelbiron]

As in pyro-ppl/numpyro#1952

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.66%. Comparing base (f6c1b00) to head (e1961e6).
Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main       #4   +/-   ##
=======================================
  Coverage   96.66%   96.66%           
=======================================
  Files          13       13           
  Lines         510      510           
=======================================
  Hits          493      493           
  Misses         17       17           

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