ExpertGate Nan loss using multi-head classifier

In my current experiments, I am trying to set up a different dataset ( with different number of classes) as separate tasks/experiences. While trying to train on ExpertGate, I am not sure if I am doing the data processing correctly as I am getting `Nan` loss repeatedly.

Here's my dataprocessing code:

```python
for idx, (task_name, num_classes) in enumerate(zip(task_list, classes_per_task)):
print("task: ", task_name)
_, _, _, train_dataset, test_dataset = get_dataloaders(task_name, 0.8, 16)
train_avalanche_dataset = _make_taskaware_classification_dataset(train_dataset)

test_avalanche_dataset = _make_taskaware_classification_dataset(test_dataset)
train_dataset_list.append(train_avalanche_dataset)
test_dataset_list.append(test_avalanche_dataset)

ncbm = nc_benchmark(
train_dataset_list,
test_dataset_list,
n_experiences = 100,
task_labels=True,
shuffle=False,
class_ids_from_zero_in_each_exp=True,
one_dataset_per_exp=True,
train_transform=None,
eval_transform=None,
)
```

This is how I am initializing the model and strategy:

```python
model = ExpertGate(shape=(3, 224, 224),device=device)
model.expert.classifier[6] = MultiHeadClassifier(4096)
cl_strategy = ExpertGateStrategy(model = model, optimizer = optimizer,
train_mb_size=256,
eval_mb_size=128,
train_epochs=2,
ae_lr=1e-3,
device = device
)
```

Here's a sample (on-going) training log:

```bash
Device: cuda:0
Starting experiment...
Start training on experience 0
-- >> Start of training phase << --

TRAINING NEW AUTOENCODER
-- >> Start of training phase << --
15106it [09:10, 27.44it/s]
Epoch 0 ended.
Loss_Epoch/train_phase/train_stream/Task000 = 135.1136
100%|██████████| 5241/5241 [00:55<00:00, 95.15it/s]
Epoch 1 ended.
Loss_Epoch/train_phase/train_stream/Task000 = 124.1420
-- >> End of training phase << --
FINISHED TRAINING NEW AUTOENCODER

SELECTING EXPERT
FINISHED EXPERT SELECTION

TRAINING EXPERT
100%|██████████| 21/21 [00:34<00:00, 1.65s/it]
Epoch 0 ended.
Loss_Epoch/train_phase/train_stream/Task000 = nan
Top1_Acc_Epoch/train_phase/train_stream/Task000 = 0.0130
100%|██████████| 21/21 [00:34<00:00, 1.64s/it]
Epoch 1 ended.
Loss_Epoch/train_phase/train_stream/Task000 = nan
Top1_Acc_Epoch/train_phase/train_stream/Task000 = 0.0036
-- >> End of training phase << --
-- >> Start of eval phase << --
-- Starting eval on experience 0 (Task 0) from test stream --
100%|██████████| 11/11 [00:09<00:00, 1.21it/s]
> Eval on experience 0 (Task 0) from test stream ended.
Loss_Exp/eval_phase/test_stream/Task000/Exp000 = nan
Top1_Acc_Exp/eval_phase/test_stream/Task000/Exp000 = 0.0061
-- >> End of eval phase << --
Loss_Stream/eval_phase/test_stream/Task000 = nan
Top1_Acc_Stream/eval_phase/test_stream/Task000 = 0.0061
End training on experience 0
Training time: 191.16707921028137
Computing accuracy on the test set
Start training on experience 1
-- >> Start of training phase << --

TRAINING NEW AUTOENCODER
-- >> Start of training phase << --
30%|███ | 3756/12490 [01:01<02:16, 64.00it/s]
1
​
1
```

Would love to hear any pointers on this?
In general, what is the best way to set up dataloaders for my particular setting?