Fix casting step outputs to device

[gsarti]

Description

Ensures the outcome of the attribute_step method is always on cpu to minimize GPU memory allocation, while also preventing device issues raised in #251

[gsarti]

@xuan25 could you try out whether this works for you on CUDA? The current version we merged from your PR causes some issues when out.show() is called because tensors are never cast back to CPU and the call tensor.numpy() in the viz method fails. This should fix it but preserve CUDA compatibility. Thanks in advance!

[xuan25]

Hi @gsarti, thanks for pointing it out. I can confirm all 50 tests have been passed on my side with this branch.

I also think whether it is more reasonable to only move the tensor to the CPU when show or other operations that have to be done on the CPU have been invoked. Otherwise, the tensor can be left on its original device (same as the model) to avoid potentially unnecessary overhead.

[gsarti]

Thanks @xuan25, wil merge then! I think that if the goal is to make the methods as accessible as possible, saving a bit of GPU RAM by moving tensors back to CPU should be the preferred approach as it could enable people to use bigger batch sizes or models. I think it's worth keeping as-is in the current state, but for the future an option would be to have a flag to control this behavior, while still keeping the memory-saving one as default!