skorch hyper-parameters grid search fails for classification task with probability vectors as labels (pytorch)

I am trying to use the skorch library to find the best hyper parameters for my simple NN model.

My labels are a probabily vectors, that asing some probabily (gaussian) to 10 classes, for example:

y[0] = [0, 0, 0.1, 0.2, 0.5, 0.1, 0.1, 0, 0, 0]

My model is a simple FC model that outputs 10 values, which are the confidence level for each class:

class Predictor(torch.nn.Module):
    def __init__(self, input_size, layer_size, layer_num, dropout=0.2):
        super().__init__()
        self.hidden = torch.nn.ModuleList()
        self.fc_in = torch.nn.Linear(input_size, layer_size)
        for i in range(layers_num):
            self.hidden.append(torch.nn.Linear(layer_size, layer_size))
        self.fc_out = torch.nn.Linear(layer_size, 10)
        self.drop = torch.nn.Dropout(p=dropout)

    def forward(self, x):
        x = F.relu(self.drop(self.fc_in(x)))
        for layer in self.hidden:
            x = F.relu(self.drop(layer(x)))
        x = F.relu(self.fc_out(x))
        return x

And I am using the skorch NeuralNetClassifier for the grid search:

data = CustomDataset("train_data.csv")
dataloader = DataLoader(data, batch_size=len(data), shuffle=True, num_workers=4)
batch, labels = next(iter(dataloader))

val_data = CustomDataset("validation_data.csv")
val_dataloader = DataLoader(val_data, batch_size=len(val_data), shuffle=True, num_workers=4)
batch_validation, labels_validation = next(iter(val_dataloader))

batch = torch.tensor(batch, dtype=torch.float32)
labels = torch.tensor(labels, dtype=torch.float32)

model = NeuralNetClassifier(
    Predictor,
    criterion=nn.CrossEntropyLoss,
    optimizer=optim.Adam,
    train_split=predefined_split(val_data),
    module__input_size=1148,
    module__layer_size=128,
    module__layers_num=3,
    verbose=True
)

# define the grid search parameters
param_grid = {
    'batch_size': [10, 20],
    'max_epochs': [10, 50]
}
grid = GridSearchCV(estimator=model, param_grid=param_grid, n_jobs=3, cv=3)
grid_result = grid.fit(batch, labels)

And the problem is that I am always getting this error:

ValueError: Classification metrics can't handle a mix of continuous-multioutput and binary targets

Any ideas?

I did try to apply softmax instead of ReLU in the output of the model and it didn't work. Also tried to use my custom scoring function and nothing changes.

NeuralNetClassifier is expecting binary targets (i.e., single-label classification) while your labels are probability vectors for multi-label classification. Use softmax to convert to probabilities:

import torch.nn.functional as F

class Predictor(torch.nn.Module):
    def __init__(self, input_size, layer_size, layers_num, dropout=0.2):
        super().__init__()
        self.hidden = torch.nn.ModuleList()
        self.fc_in = torch.nn.Linear(input_size, layer_size)
        for _ in range(layers_num):
            self.hidden.append(torch.nn.Linear(layer_size, layer_size))
        self.fc_out = torch.nn.Linear(layer_size, 10)
        self.drop = torch.nn.Dropout(p=dropout)

    def forward(self, x):
        x = F.relu(self.drop(self.fc_in(x)))
        for layer in self.hidden:
            x = F.relu(self.drop(layer(x)))
        x = self.fc_out(x)
        return F.softmax(x, dim=1)  # Apply softmax to convert to probabilities

Then reshape your labels to match the expected format for multi-label classification like this:

labels = torch.argmax(labels, dim=1)