pytorch CNN 모델 output shape 관련 질문

CIFAR10 데이터셋을 이용해서 cnn을 구현하고 training 실습을 하는데 계속 모델의 output shape과 관련해서 CrossEntropyLoss를 이용해서 loss를 계산하는 부분에서

ValueError: Expected input batch_size (1) to match target batch_size (4).

와 같은 에러가 발생합니다.

input과 labels의 배치 사이즈에 따라 마지막 output의 shape이 torch.Size([4, 10])와 같이 나와야 정상적으로 loss를 계산할 수 있을텐데 계속해서 torch.Size([1, 10])의 shape이 나옵니다.

혹시몰라 좀 더 간단한 다른 모델로 구현해보니 정상적으로 batch_size가 일치하는데 해당 모델을 이용할 때만 오류가 발생합니다.

batch_size를 바꿔가면서 테스트해보니 모델에 어떤 문제로 인해 계속 output의 batch_size가 설정한 batch_size의 1/4로 나오게 되는 것 같은데 모델에서 어떤 부분이 문제인지 못 찾겠습니다.

CNN 모델은 아래와 같이 구현했습니다. 모델의 output shape이 정상적으로 torch.Size([batch_size, 10])과 같이 나올 수 있으려면 어떻게 해결해야 할까요?

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        # Layer#1
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3)

        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.dropout = nn.Dropout(0.5)

        # Layer#2
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3)

        # Layer#3
        self.conv4 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3)

        # Layer#4
        self.fc1 = nn.Linear(4*4*256, 512)
        # Layer#5
        self.fc2 = nn.Linear(512, 512)
        # Layer#6
        self.fc3 = nn.Linear(512, 10)



    def forward(self, x):

        # Layer#1
        x = self.conv1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.pool(x)

        # Layer#2
        x = self.conv3(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.pool(x)

        # Layer#3
        x = self.conv4(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.pool(x)

        # Layer#4
        x = x.view(-1, 4*4*256)
        x = self.fc1(x)
        x = F.relu(x)

        # Layer#5
        x = self.fc2(x)
        x = F.relu(x)

        # Layer#6
        x = self.fc3(x)
        x = F.softmax(x)

        return x