Given the skeleton of the neutral network of mnasnet1_0 and the evaluation function as below , I checked that some of my prediction outputs are even not a digital integer between 0 and 9 in MNIST dataset, hence, I got a very low accuracy for the model. May I know any hints for me to correct my model? Or, how to evaluate the performance of mnasnet 1.0 model with mnist dataset?

Thanks.

`class MnasMulti(nn.Module):
    
    def __init__(self, alpha=1.0):
        super(MnasMulti, self).__init__()
        depths = _get_depths(alpha)
        if alpha == 1.0:
            MNASNet = torchvision.models.mnasnet1_0(pretrained=True, progress=True)
        else:
            MNASNet = torchvision.models.MNASNet(alpha=alpha)

        self.conv0 = nn.Sequential(
            MNASNet.layers._modules[ 0 ],
            MNASNet.layers._modules[ 1 ],
            MNASNet.layers._modules[ 2 ],
            MNASNet.layers._modules[ 3 ],
            MNASNet.layers._modules[ 4 ],
            MNASNet.layers._modules[ 5 ],
            MNASNet.layers._modules[ 6 ],
            MNASNet.layers._modules[ 7 ],
            MNASNet.layers._modules[ 8 ],
        )

        self.conv1 = MNASNet.layers._modules[ 9 ]
        self.conv2 = MNASNet.layers._modules[ 10 ]

        self.out1 = nn.Conv2d(depths[4], depths[4], 1, bias=False)
        final_chs = depths[4]
        self.inner1 = nn.Conv2d(depths[3], final_chs, 1, bias=True)
        self.inner2 = nn.Conv2d(depths[2], final_chs, 1, bias=True)

        self.out2 = nn.Conv2d(final_chs, depths[3], 3, padding=1, bias=False)
        self.out3 = nn.Conv2d(final_chs, depths[2], 3, padding=1, bias=False)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)
        # conv3 = self.conv3(conv2)

        intra_feat = conv2
        outputs = []
        out = self.out1(intra_feat)
        outputs.append(out)

        intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner1(conv1)
        out = self.out2(intra_feat)
        outputs.append(out)

        intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner2(conv0)
        out = self.out3(intra_feat)
        outputs.append(out)

        return outputs[::-1]

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    acc = 100. * correct / len(test_loader.dataset)
    print( 
Test set: Accuracy: {}/{} ({:.2f}%)
 .format(correct, len(test_loader.dataset), acc))

    return`

====================

I had reshaped my pred in order to fit in the size of the target, but when i printed out the values of pred, some of the outputs are not a single digits :(