Свертка 1x1 как слой классификации в Pytorch

Я пытаюсь классифицировать патчи изображений по 10 различным категориям, используя нейронную сеть. Моя идея (позаимствовано из эта статья предназначена для использования первых 5 слоев предварительно обученной сети VGG и применения к этому кодировщику свертки 1 x 1. Итак, учитывая первые 5 слоев сети VGG:

Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

Как я могу добавить слой для классификации каждой активации в n_classes? Я пытаюсь выполнить свертку 1x1 от 512 (количество фильтров в последнем слое) до n_classes, что-то вроде этого:

nn.Sequential(
    nn.Conv2d(in_channels=512, out_channels=n_classes, kernel_size=1),
    nn.Softmax(dim=1)
    )

Эта сеть дает (batch_size, n_classes H // 32, W // 32) выходных данных. Однако все выходные данные для каждого класса равны для данного изображения, я имею в виду, что для данной партии B и класса C все значения одинаковы.

Я пробовал тренироваться с CrossEntropyLoss и даже конвертировать в логиты и применять BCELoss, но безуспешно. Я также пробовал другие сети в своем цикле поездов, и этого не происходит.

Любая подсказка?

Спасибо