Problem in 'Captum'

sampa: <p>Captum is working without any error. But every time It is predicted as first class. Here is the latest snapshot of my code</p>
<pre><code class="lang-auto">
import os
from captum.insights import AttributionVisualizer, Batch
from captum.insights.features import ImageFeature
from torchvision import datasets, transforms, models
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
test_transforms = transforms.Compose([
                    transforms.Resize((224,224)),
                    transforms.ToTensor(),
                    #transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010))
                    ])

def get_classes():
    classes = [
        "B",
        "M" 
    ]
    return classes


def get_pretrained_model():
    net= models.resnet18(pretrained=True)
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 2)
    pt_path = os.path.abspath(
        os.path.dirname(__file__) + "/models/ResnetBW.pt"
    )
    net.load_state_dict(torch.load(pt_path))
    return net


def baseline_func(input):
    return input * 0


def formatted_data_iter():
    dataset = datasets.ImageFolder(
        root="data/test", transform=test_transforms
    )
    dataloader = iter(
        torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=False, num_workers=2)
    )
    while True:
        images, labels = next(dataloader)
        yield Batch(inputs=images, labels=labels)


if __name__ == "__main__":
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    model = get_pretrained_model()
    visualizer = AttributionVisualizer(
        models=[model],
        score_func=lambda o: torch.nn.functional.softmax(o, 1),
        classes=get_classes(),
        features=[
            ImageFeature(
                "Photo",
                baseline_transforms=[baseline_func],
                input_transforms=[normalize],
            )
        ],
        dataset=formatted_data_iter(),
    )

    visualizer.render(debug=True)

</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/3/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6.png" data-download-href="https://discuss.pytorch.org/uploads/default/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6" title="Capture"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6_2_517x346.png" alt="Capture" data-base62-sha1="9FDvC1VE7xkDSh7t7mZcpewou7c" width="517" height="346" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6_2_517x346.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6_2_775x519.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/3/43c87a09c46e52ca19b5bd75c5a1cdcdece4c8f6_2_1034x692.png 2x" data-dominant-color="E9E9E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Capture</span><span class="informations">1277×856 240 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Are you getting other results, if you don’t use captum?<br>
Could it be that your model just overfits class0?</p>

sampa: <p>No…There is no problem in model. I am testing the captum with the images which are predicted true class with this model.</p>

sampa: <p>Instead of class ‘B’ and "M’. If I use ‘N’ and ‘M’. The model is predicted all image as ‘M’ class.</p>