TST_captum.py

from captum.attr import IntegratedGradients
import pickle
import torch

with open("./tst_save.pkl", "rb") as f:
  trained_model = pickle.load(f)
  
 trained_model.eval()

X = torch.load("./X_test.pt")
X.requires_grad = True

attributor = IntegratedGradients(trained_model)

attributions = attributor.attribute(X, target=0)

assert attributions.shape == X.shape

torch.save(attributions, "./attributions.pt")