How can Fig. 6 be reproduced?? (Effective Receptive Field)

According to https://arxiv.org/pdf/1701.04128.pdf:

we place a gradient signal of 1 at the center of the output plane and 0 everywhere else, and then back-propagate this gradient through the network to get input gradients.

Something like this should work (not sure about dimension order):

def erf(image, model):
  b, h, w, c = image.shape
  image = image.requires_grad(True)
  output = model(image)
  gradient = torch.zeros_like(output)
  gradient[0, h//2, w//2, :] = 1
  gradient.requires_grad = True
  output.backwards(gradient = gradient)

  return image.grad.detach()

NVlabs / SegFormer

How can Fig. 6 be reproduced?? (Effective Receptive Field) #123