IncrementalClassifier model not adapting correctly

[PerezLucas2903]

I'm currently running some tests to learn how to properly use Avalanche in my research. In this test, I try to train a simple MLP with SplitMNIST as follows:

benchmark = SplitMNIST(n_experiences=5, seed=1)

shape = np.array(benchmark.train_stream[0].dataset[0][0].shape)
in_features = shape.prod()
baseline_model = IncrementalClassifier(in_features = in_features, initial_out_features=2, masking=True)
optimizer = SGD(baseline_model.parameters(), lr=0.001, momentum=0.9)
criterion = CrossEntropyLoss()

params = {'batch_size': 512,
      'num_workers': 2}

for experience in benchmark.train_stream:

    print(experience.classes_in_this_experience)

    baseline_model.adaptation(experience)
    print(baseline_model)

    data_tensor = torch.empty((0, 784), dtype=torch.float32)
    data_label = []

    for data in experience.dataset:
        data_ravel = torch.ravel(data[0])
        data_reshaped = torch.reshape(data_ravel, (-1, 784))
        data_tensor = torch.cat((data_tensor, data_reshaped), 0)
        data_label.append(data[1])

    custom_ds = CustomDataset(data_tensor, data_label)
    train_loader = DataLoader(custom_ds, **params)


    for X, y in train_loader:
        baseline_model.train()
        optimizer.zero_grad()
        output = baseline_model.forward(X)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()_

My intention is to acess the model's parameters as I train each experience. However, when I print my model, the number of classification layers is wrong on all experiences. It goes as follows:

[5, 6]
IncrementalClassifier(
(classifier): Linear(in_features=784, out_features=7, bias=True)
)
[1, 2]
IncrementalClassifier(
(classifier): Linear(in_features=784, out_features=7, bias=True)
)
[0, 8]
IncrementalClassifier(
(classifier): Linear(in_features=784, out_features=9, bias=True)
)
[9, 3]
IncrementalClassifier(
(classifier): Linear(in_features=784, out_features=10, bias=True)
)
[4, 7]
IncrementalClassifier(
(classifier): Linear(in_features=784, out_features=10, bias=True)
)

When it should go 2 by 2 all the way to 10.
What am I doing wrong here?
(This is my first time writing in forums so excuse my poor english)

[AntonioCarta]

Labels are not ordered. So, when the model sees [5, 6], it expands the classifier to account for 7 classes (0 to 6). Unseen clases (0, 1, 2, 3) are masked by default.

[PerezLucas2903]

I see, thanks! Is there any way around this so I can test it multiple times with random digits per experience without expanding too much of my network?

[AntonioCarta]

you can order the labels yourself and set [0, 1, 2, ...] as the class order in the class-incremental benchmark generator.