Solved by overflow_question in post #2 
              
              
                I resolved the issue with LayerIntegratedGradients. 
Here is the link to read more to know other possible solutions. Captum · Model Interpretability for PyTorch 
Sample code 
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from captum.attr import IntegratedGradi…