Generalize technique (first steps)

[navneeth]

Hi
Thank you for devising and sharing a nice technique for explainability of CNNs.
I would like generalize this code to other datasets and architectures.
To this effect I made a few small changes to it. Specifically

• remove imagenet specific code
• replace .cuda calls with .device calls.

Could you review and merge the PR. Bulding on this PR, I would like to add more general datasets, architectures to it.

thanks.

[navneeth]

is this feature desirable for this method ? Do check it out.