[BUG] Colab demo for saliency maps: Gradient images are blank about flashtorch

[drscotthawley]

Perhaps something important has changed on Colab since these examples were created: The other Colab demo on activation minimization also fails, although in that case it gives a RuntimeError. (RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!)

These problems only seem to affect the Colab notebooks. When I cloned the repo and ran locally using the non-colab notebooks, I was able to see the saliency maps ok.

[mpierrau]

I tried debugging a bit but only managed to figure out that the problem lies in that the gradient returned from AlexNet in the backpropagation does not correspond to the expected dimensions, so self.gradients is never updated because it fails the if statement in _record_gradients:

grad_in[0].shape: torch.Size([1, 64, 55, 55])
self.gradients.shape: torch.Size([1, 3, 224, 224])

When running with CPU or TPU on colab it works fine!