import torch
from avalanche.benchmarks import GenericCLScenario, dataset_benchmark
from avalanche.benchmarks.utils import AvalancheTensorDataset
from torchvision import transforms

import ctrl


def CTrL(stream_name):
    stream = ctrl.get_stream(stream_name)

    exps = [[], [], []]
    norms = []
    for t in stream:
        for split, exp in zip(t.datasets, exps):
            samples, labels = split.tensors
            # samples -= torch.tensor(t.statistics['mean']).view(1, 3, 1, 1)
            # samples /= torch.tensor(t.statistics['std']).view(1, 3, 1, 1)

            task_labels = [t.id] * samples.size(0)
            dataset = AvalancheTensorDataset(samples, labels.squeeze(1),
                                             task_labels=task_labels)
            exp.append(dataset)
        norms.append(transforms.Normalize(t.statistics['mean'],
                                          t.statistics['std']))
    return dataset_benchmark(
        train_datasets=exps[0],
        test_datasets=exps[2],
        other_streams_datasets=dict(valid=exps[1]),
    )

@@ -254,16 +254,54 @@
 _name)%0A%0A
+    # Train, val and test experiences%0A
     exps
@@ -320,23 +320,8 @@
 %5B%5D%5D%0A
-    norms = %5B%5D%0A
     
@@ -349,169 +349,165 @@
     
-for split, exp in zip(t.datasets, exps):%0A            samples, labels = split.tensors%0A            # samples -= torch.tensor(t.statistics%5B'mean'%5D).view(1, 3, 1, 1)
+trans = transforms.Normalize(t.statistics%5B'mean'%5D,%0A                                     t.statistics%5B'std'%5D)%0A        for split, exp in zip(t.datasets, exps):
 %0A   
@@ -518,18 +518,16 @@
     
- #
  samples
  /= 
@@ -526,63 +526,32 @@
 ples
- /= torch.tensor(t.statistics%5B'std'%5D).view(1, 3, 1, 1)%0A
+, labels = split.tensors
 %0A   
@@ -739,17 +739,17 @@
 k_labels
-)
+,
 %0A       
@@ -757,154 +757,88 @@
     
-exp.append(dataset)%0A        norms.append(transforms.Normalize(t.statistics%5B'mean'%5D,%0A                                          t.statistics%5B'std'%5D)
+                                 transform=trans)%0A            exp.append(dataset
 )%0A  

Add normalization to each task