The size of tensor a (512) must match the size of tensor b (2048) at non-singleton dimension 1

I'm trying to use shap to improve the explainability of my model. This is the code-

import shap 
batch = next(iter(test_dl))
images, _ = batch

background = images[:100].to(device)
test_images = images[100:105].to(device)

e = shap.DeepExplainer(model, background)
shap_values = e.shap_values(test_images)

shap_numpy = [np.swapaxes(np.swapaxes(s, 1, -1), 1, 2) for s in shap_values]
test_numpy = np.swapaxes(np.swapaxes(test_images.cpu().numpy(), 1, -1), 1, 2)
shap.image_plot(shap_numpy, -test_numpy)

When I run this, I get the error - RuntimeError: The size of tensor a (512) must match the size of tensor b (2048) at non-singleton dimension 1

How do I solve this? I am working on a binary classification problem and I'm using ResNet50I'm trying to use shap to improve the explainability of my model. This is the code-

class PredsModel(ImageClassificationBase):
    def __init__(self, num_classes, pretrained=True):
        super().__init__()
        # Use a pretrained model
        self.network = models.resnet50 (pretrained=pretrained)
        # Replace last layer
        self.network.fc = nn.Linear(self.network.fc.in_features, num_classes)

    def forward(self, xb):
        return self.network(xb)

The full code is available here - https://colab.research.google.com/drive/1gQO_RddY0aBYtTQ2HTDcP6PXVEsJYuJL?usp=sharing

Thanks