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

Question

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.