Fix test batches and make weighted loss general to device

[dfulu]

• Fix batches for datapipes 3.1.5
• Minor fix to make sure the weighted loss will work regardless of the device used

[codecov[bot]]

Codecov Report

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

Comparison is base (3b74b1d) 57.62% compared to head (b2d0e30) 57.62%.

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #117 +/- ## ======================================= Coverage 57.62% 57.62% ======================================= Files 26 26 Lines 1706 1706 ======================================= Hits 983 983 Misses 723 723 ```

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