fix: move Exa.TrkX models to device explicitely

[benjaminhuth]

It seems like the Exa.TrkX models are not moved to device in some cases, this adds a few lines that do this explicitely.

[codecov]

Codecov Report

Merging #1664 (f80fe85) into main (c2de690) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #1664   +/-   ##
=======================================
  Coverage   48.46%   48.46%           
=======================================
  Files         384      384           
  Lines       21071    21071           
  Branches     9702     9702           
=======================================
  Hits        10213    10213           
  Misses       4130     4130           
  Partials     6728     6728           

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[github-actions]

📊 Physics performance monitoring for f80fe85

🟥 ERROR The result has missing elements!
This is likely a physmon job failure

Full report
CKF: ❌ seeded, ❌ truth smeared, ❌ truth estimated
IVF: ❌ seeded, ❌ truth smeared, ❌ truth estimated
❌ Ambiguity resolution
❌ Truth tracking

Vertexing ❌

❌ Vertexing vs. mu

❌ IVF seeded

❌ IVF truth_smeared

❌ IVF truth_estimated

CKF ❌

❌ CKF seeded

❌ CKF truth_smeared

❌ CKF truth_estimated

Ambiguity resolution ❌

❌ seeded

Truth tracking ❌

❌ Truth tracking