Test failed on CUDA enabled platform

[xuan25]

🐛 Bug Report

Test failed on CUDA enabled platform.

Output is shown below.

=============================================================================================================================== FAILURES ===============================================================================================================================
_____________________________________________________________________________________________________________________ test_attribute_slice_seq2seq _____________________________________________________________________________________________________________________
[gw4] linux -- Python 3.10.12 <hide>/inseq/.venv/bin/python
tests/models/test_huggingface_model.py:256: in test_attribute_slice_seq2seq
    out = saliency_mt_model.attribute(
inseq/models/attribution_model.py:445: in attribute
    attribution_outputs = attribution_method.prepare_and_attribute(
inseq/attr/attribution_decorators.py:72: in batched_wrapper
    out = f(self, *args, **kwargs)
inseq/attr/feat/feature_attribution.py:237: in prepare_and_attribute
    attribution_output = self.attribute(
inseq/attr/feat/feature_attribution.py:431: in attribute
    step_output = self.filtered_attribute_step(
inseq/attr/feat/feature_attribution.py:596: in filtered_attribute_step
    step_output.remap_from_filtered(target_attention_mask, orig_batch)
inseq/data/attribution.py:431: in remap_from_filtered
    self.source_attributions = remap_from_filtered(
.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py:115: in decorate_context
    return func(*args, **kwargs)
inseq/utils/torch_utils.py:35: in remap_from_filtered
    return new_source.scatter(0, index, filtered)
E   RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA_scatter_src)
======================================================================================================================= short test summary info ========================================================================================================================
FAILED tests/models/test_huggingface_model.py::test_attribute_slice_seq2seq - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA_scatter_src)
======================================================================================================== 1 failed, 49 passed, 37 warnings in 340.47s (0:05:40) =========================================================================================================
make: *** [Makefile:101: fast-test] Error 1

🔬 How To Reproduce

1. Clone the repo
2. Setup the dev environment
3. Run make fast-test

Code sample

N/A

Environment

• OS: Windows Subsystem Linux (Ubuntu 22.04.2 LTS)
• Python version: Python 3.10.12
• Inseq version: 0.6.0.dev0

Expected behavior

All tests passed

Additional context

CUDA Version: 12.3