Need advice after training data

So i have a network for image classification, for now there are just 2 classes, here is my network

class ActionNet(Module):
    def __init__(self, num_class=4):
        super(ActionNet, self).__init__()
        
        self.cnn_layer = Sequential(
            #conv1
            Conv2d(in_channels=1, out_channels=32, kernel_size=1, bias=False),
            BatchNorm2d(32),
            PReLU(num_parameters=32),
            MaxPool2d(kernel_size=3),
            #conv2
            Conv2d(in_channels=32, out_channels=64, kernel_size=1, bias=False),
            BatchNorm2d(64),
            PReLU(num_parameters=64),
            MaxPool2d(kernel_size=3),
            #flatten
            Flatten(),
            Linear(576, 128),
            BatchNorm1d(128),
            ReLU(inplace=True),
            Dropout(0.5),
            Linear(128, num_class)
        )
    
    def forward(self, x):
        x = self.cnn_layer(x)
        return x

then i train with 600 images without augmentation , 300 images each class. i split the data to 8:2 train:test. the result is good enough i get 87% accuracy for test data but the training loss and validation loss is far.
i trained data with cross entropy loss function with learning rate 0.001 and 50 epoch training.

here is my example training data


class: A


class: B

when i try to test with new data like this image (it should be non of those class):

the result of output model is
tensor([[1.5626, 0.2562]], device='cuda:0', grad_fn=<AddmmBackward>)
which is 1.56 for class A and 0.25 class B, of course the result is class A but it should be non of them.

any advice for this would be appreciate, thank you

In Image classification,the model predicts high value tensor for most likely class. So there’s no scope for a model to predict nothing as it predicts the most probable class associated to the image. So the model can’t give the result as no class associated to it. Whichever tensor has higher value will be the most probable class.

any solution? add more class or what?

You can add 1 more label
For eg :
Label 0 - No class
Label 1 - Class A
Label 2 - Class B

do i need to add dataset for no class too? or just train with same data

Yeah you have to add data for no class too