CUDA/CPU transfer error

[rotobash]

Hi, I was interested in visualizing a CNN I created. I was following the example provided but when I try to use my model instead of a pretrained model from torchvision I get:

File "cnn_layer_visualization.py", line 82, in <module>
    gradients = backprop.calculate_gradients(input_, target_class)
  File "C:\Users\Mike\Anaconda3\lib\site-packages\flashtorch\saliency\backprop.py", line 106, in calculate_gradients
    output.backward(gradient=target)
  File "C:\Users\Mike\Anaconda3\lib\site-packages\torch\tensor.py", line 107, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\Users\Mike\Anaconda3\lib\site-packages\torch\autograd\__init__.py", line 93, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: invalid gradient at index 0 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

I looked through the source code a bit and saw that Backprop moves the network to cuda (if available) regardless if I move it there myself or not. Is this possibly a bug? I've had a bit of trouble with CUDA in the past so it might be a problem on my end.

Thanks for your time, if there's any other information I can provide please let me know. Cheers.

[MisaOgura]

Hi @rotobash thanks for submitting an issue.

I have reproduced the error - it is due to target tensor not being transferred to GPU.

Also, perhaps a better API design would be to allow users to explicitly set a device to use, rather than automatically using GPU if available.

I'm working on it and will open a PR soon :)

[MisaOgura]

Hi again @rotobashm, the fix has been merged and the updated package is published as v0.0.8 (release note).

Let me know whether the issue gets solved, so I can close this :)

[rotobash]

Hi @MisaOgura, so sorry for the delay I hadn't had a chance to look at this till now. This seems to solve the cuda transfer issue for the one model I tried, I'll try some other models when I get a chance and let you know if I run into any issues. Thanks for the fix!

[MisaOgura]

Brilliant @rotobash - I would love to see flashtorch being used in other projects, and would love to list them on README in the future. So please do let me know how yours goes :)