I am attempting to build a convulational neural network from scratch for object recognition. My input image size is (3, 256, 256)
Here’s my architecture for CNN.
class NutSnackClassication(MultiLabelImageClassificationBase):
    def __init__(self):
        super().__init__()
        
        self.network = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            nn.Conv2d(32, 32, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(2,2),
            
            nn.Conv2d(32, 64, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(128, 256, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1), 
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            
            nn.Flatten(),
            nn.Linear(7*7*256, 512),
            nn.ReLU(),
            nn.Linear(512, 258),
            nn.LogSoftmax(dim = 1),
        )
    
    def forward(self, xb):
        return self.network(xb)

I have to use Weighted accuracy of all classes so I chose to use F1 Score as metric to evaluate my models performance.
def F_score(output, label, threshold=0.5, beta=1):
    prob = output > threshold
    label = label > threshold
    print(prob.shape, "---", label.shape)
    TP = (prob & label).sum(1).float()
    TN = ((~prob) & (~label)).sum(1).float()
    FP = (prob & (~label)).sum(1).float()
    FN = ((~prob) & label).sum(1).float()

    precision = torch.mean(TP / (TP + FP + 1e-12))
    recall = torch.mean(TP / (TP + FN + 1e-12))
    F2 = (1 + beta**2) * precision * recall / (beta**2 * precision + recall + 1e-12)
    return F2.mean(0)

Specifically I am using this function below:
    # validation step
    def validation_step(self, batch):
        images, targets = batch
        out = self(images)
        loss = F.nll_loss(out, targets)
        score = F_score(out, targets)
        return {'val_loss': loss.detach(), 'val_score': score.detach() }

But when I execute this evaluate(model, val_dl)  I get an error saying:
torch.Size([128, 258]) --- torch.Size([128])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)


RuntimeError: The size of tensor a (258) must match the size of tensor b (128) at non-singleton dimension 1

I dont understand how to solve this. Is my F1_Score() rightly implemented? How do I approach a solution to this?