use of custom model throws an error about flashtorch HOT 2 CLOSED

Hi @SarfarazHabib,

Thanks for your question.

The reason why it's giving an error is that currently the visualize method is passing the output of the gradient ascent to matplotlib which expects the input to be an RBG image.

It is on my todo list to make the API more generalisable to various custom inputs and not just RGB or grayscale images.

As a quick solution meanwhile, I suggest below:

1. Delete this line so the conditional becomes just

if isinstance(module, nn.modules.conv.Conv2d):

2. Use optimize api, rather than visualize, which returns raw gradients accumulated over the number of iterations as a list

3. Plot the gradients/channels of your interest

This is untested, but hope this helps - let me know how it goes.

from flashtorch.

@SarfarazHabib please feel free to reopen the issue or open a new one if required.

from flashtorch.