RuntimeError: expected device cpu but got device cuda:0 when training and visualizing model on IMDB about captum

[NarineK]

@finiteautomata, we fixed it on the master and haven't released a new version since.
Here is the fix: https://github.com/pytorch/captum/blob/master/captum/attr/_utils/attribution.py#L263

[NarineK]

@finiteautomata, does it answer to your question ? Can we close the task ?

[finiteautomata]

Yes, sorry! I've just tried this and it works. The only thing to mention (as already pointed out in #192) is that I had to clone the repo and install it with pip install -e .

Thanks @NarineK for your help