diff --git a/higher/config/bad_tf_config.json b/higher/config/bad_tf_config.json
new file mode 100644
index 0000000..9f25c08
--- /dev/null
+++ b/higher/config/bad_tf_config.json
@@ -0,0 +1,193 @@
+[
+ {
+ "name": "Identity",
+ "function": "identity"
+ },
+ {
+ "name": "FlipUD",
+ "function": "flip",
+ "param": {
+   "axis": "Y"
+  }
+ },
+ {
+ "name": "FlipLR",
+ "function": "flip",
+ "param": {
+   "axis": "X"
+  }
+ },
+ {
+ "name": "Rotate",
+ "function": "rotate",
+ "param": {
+   "min": null,
+   "max": 30,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateX",
+ "function": "translate",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateY",
+ "function": "translate",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearX",
+ "function": "shear",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearY",
+ "function": "shear",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Contrast",
+ "function": "contrast",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Color",
+ "function": "color",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Brightness",
+ "function": "brightness",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Sharpness",
+ "function": "sharpness",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Posterize",
+ "function": "posterize",
+ "param": {
+   "min": 4.0,
+   "max": 8.0,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Solarize",
+ "function": "solarize",
+ "param": {
+   "min": 0.00390625,
+   "max": 1.0,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BTranslateX",
+ "function": "translate",
+ "param": {
+   "axis": "X",
+   "min": 25,
+   "max": 30,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BTranslateX-",
+ "function": "translate",
+ "param": {
+   "axis": "X",
+   "min": -25,
+   "max": -30,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BTranslateY",
+ "function": "translate",
+ "param": {
+   "axis": "Y",
+   "min": 25,
+   "max": 30,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BTranslateY-",
+ "function": "translate",
+ "param": {
+   "axis": "Y",
+   "min": -25,
+   "max": -30,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BShearX",
+ "function": "shear",
+ "param": {
+   "axis": "X",
+   "min": 0.9,
+   "max": 1.2,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "BShearY",
+ "function": "shear",
+ "param": {
+   "axis": "Y",
+   "min": 0.9,
+   "max": 1.2,
+   "absolute": true,
+   "invScale": false
+  }
+ }
+]
diff --git a/higher/config/base_tf_config.json b/higher/config/base_tf_config.json
new file mode 100644
index 0000000..a8967ac
--- /dev/null
+++ b/higher/config/base_tf_config.json
@@ -0,0 +1,127 @@
+[
+ {
+ "name": "Identity",
+ "function": "identity"
+ },
+ {
+ "name": "FlipUD",
+ "function": "flip",
+ "param": {
+   "axis": "Y"
+  }
+ },
+ {
+ "name": "FlipLR",
+ "function": "flip",
+ "param": {
+   "axis": "X"
+  }
+ },
+ {
+ "name": "Rotate",
+ "function": "rotate",
+ "param": {
+   "min": null,
+   "max": 30,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateX",
+ "function": "translate",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateY",
+ "function": "translate",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearX",
+ "function": "shear",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearY",
+ "function": "shear",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Contrast",
+ "function": "contrast",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Color",
+ "function": "color",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Brightness",
+ "function": "brightness",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Sharpness",
+ "function": "sharpness",
+ "param": {
+   "min": 0.1,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Posterize",
+ "function": "posterize",
+ "param": {
+   "min": 4.0,
+   "max": 8.0,
+   "invScale": false
+  }
+ },
+ {
+ "name": "Solarize",
+ "function": "solarize",
+ "param": {
+   "min": 0.00390625,
+   "max": 1.0,
+   "invScale": false
+  }
+ }
+]
diff --git a/higher/config/invScale_tf_config.json b/higher/config/invScale_tf_config.json
new file mode 100644
index 0000000..09520c8
--- /dev/null
+++ b/higher/config/invScale_tf_config.json
@@ -0,0 +1,163 @@
+[
+ {
+ "name": "Identity",
+ "function": "identity"
+ },
+ {
+ "name": "FlipUD",
+ "function": "flip",
+ "param": {
+   "axis": "Y"
+  }
+ },
+ {
+ "name": "FlipLR",
+ "function": "flip",
+ "param": {
+   "axis": "X"
+  }
+ },
+ {
+ "name": "Rotate",
+ "function": "rotate",
+ "param": {
+   "min": null,
+   "max": 30,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateX",
+ "function": "translate",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "TranslateY",
+ "function": "translate",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.33,
+   "absolute": false,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearX",
+ "function": "shear",
+ "param": {
+   "axis": "X",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "ShearY",
+ "function": "shear",
+ "param": {
+   "axis": "Y",
+   "min": null,
+   "max": 0.3,
+   "absolute": true,
+   "invScale": false
+  }
+ },
+ {
+ "name": "+Contrast",
+ "function": "contrast",
+ "param": {
+   "min": 1.0,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "+Color",
+ "function": "color",
+ "param": {
+   "min": 1.0,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "+Brightness",
+ "function": "brightness",
+ "param": {
+   "min": 1.0,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "+Sharpness",
+ "function": "sharpness",
+ "param": {
+   "min": 1.0,
+   "max": 1.9,
+   "invScale": false
+  }
+ },
+ {
+ "name": "-Contrast",
+ "function": "contrast",
+ "param": {
+   "min": 0.1,
+   "max": 1.0,
+   "invScale": true
+  }
+ },
+ {
+ "name": "-Color",
+ "function": "color",
+ "param": {
+   "min": 0.1,
+   "max": 1.0,
+   "invScale": true
+  }
+ },
+ {
+ "name": "-Brightness",
+ "function": "brightness",
+ "param": {
+   "min": 0.1,
+   "max": 1.0,
+   "invScale": true
+  }
+ },
+ {
+ "name": "-Sharpness",
+ "function": "sharpness",
+ "param": {
+   "min": 0.1,
+   "max": 1.0,
+   "invScale": true
+  }
+ },
+ {
+ "name": "Posterize",
+ "function": "posterize",
+ "param": {
+   "min": 4.0,
+   "max": 8.0,
+   "invScale": true
+  }
+ },
+ {
+ "name": "Solarize",
+ "function": "solarize",
+ "param": {
+   "min": 0.00390625,
+   "max": 1.0,
+   "invScale": true
+  }
+ }
+]
diff --git a/higher/smart_aug/benchmark.py b/higher/smart_aug/benchmark.py
index d4d34be..0cff138 100644
--- a/higher/smart_aug/benchmark.py
+++ b/higher/smart_aug/benchmark.py
@@ -16,8 +16,11 @@ optim_param={
     },
     'Inner':{
         'optim': 'SGD',
-        'lr':1e-2, #1e-2 #1e-1 for ResNet
+        'lr':1e-1, #1e-2/1e-1 (ResNet)
         'momentum':0.9, #0.9
+        'decay':0.0005, #0.0005
+        'nesterov':True,
+        'scheduler':'cosine', #None, 'cosine', 'multiStep', 'exponential'
     }
 }
 
@@ -28,6 +31,7 @@ dataug_epoch_start=0
 nb_run= 3
 
 # Use available TF (see transformations.py)
+'''
 tf_names = [
     ## Geometric TF ##
     'Identity',
@@ -63,7 +67,10 @@ tf_names = [
     #'RandBlend'
 ]
 tf_dict = {k: TF.TF_dict[k] for k in tf_names}
-
+'''
+tf_config='../config/base_tf_config.json'
+TF_loader=TF_loader()
+tf_dict, tf_ignore_mag =TF_loader.load_TF_dict(tf_config)
 
 device = torch.device('cuda')
 
@@ -82,11 +89,11 @@ np.random.seed(0)
 if __name__ == "__main__":
 
     ### Benchmark ###
-    '''
+    #'''
     n_inner_iter = 1
     dist_mix = [0.5]#[0.5, 1.0]
     N_seq_TF= [3, 4]
-    mag_setup = [(True, True), (False, False)] #(FxSh, Independant)
+    mag_setup = [(False, False)] #[(True, True), (False, False)] #(FxSh, Independant)
 
     for model_type in model_list.keys():
             for model_name in model_list[model_type]:
@@ -100,7 +107,7 @@ if __name__ == "__main__":
                                 torch.cuda.reset_max_memory_cached() #reset_peak_stats
                                 t0 = time.perf_counter()
 
-                                model = getattr(model_type, model_name)(pretrained=False)
+                                model = getattr(model_type, model_name)(pretrained=False, num_classes=len(dl_train.dataset.classes))
 
                                 model = Higher_model(model, model_name) #run_dist_dataugV3
                                 if n_inner_iter!=0:
@@ -137,6 +144,7 @@ if __name__ == "__main__":
                                     'Optimizer': optim_param, 
                                     "Device": device_name, 
                                     "Memory": [max_allocated, max_cached], 
+                                    "TF_config": tf_config,
                                     "Param_names": aug_model.TF_names(), 
                                     "Log": log}
                                 print(str(aug_model),": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
@@ -150,9 +158,9 @@ if __name__ == "__main__":
 
                                 print('Execution Time : %.00f '%(exec_time))
                                 print('-'*9)
-    '''
-    ### Benchmark - RandAugment/Vanilla ###
     #'''
+    ### Benchmark - RandAugment/Vanilla ###
+    '''
     for model_type in model_list.keys():
             for model_name in model_list[model_type]:
                 for run in range(nb_run):
@@ -160,7 +168,7 @@ if __name__ == "__main__":
                     torch.cuda.reset_max_memory_cached() #reset_peak_stats
                     t0 = time.perf_counter()
 
-                    model = getattr(model_type, model_name)(pretrained=False).to(device)
+                    model = getattr(model_type, model_name)(pretrained=False, num_classes=len(dl_train.dataset.classes)).to(device)
 
                     print("{} on {} for {} epochs".format(model_name, device_name, epochs))
                     #print("RandAugment(N{}-M{:.2f})-{} on {} for {} epochs".format(rand_aug['N'],rand_aug['M'],model_name, device_name, epochs))
@@ -180,7 +188,7 @@ if __name__ == "__main__":
                         #"Rand_Aug": rand_aug, 
                         "Log": log}
                     print(model_name,": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
-                    filename = "{} epochs -{}".format(model_name,epochs, run)
+                    filename = "{}-{} epochs -{}".format(model_name,epochs, run)
                     #print("RandAugment-",model_name,": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
                     #filename = "RandAugment(N{}-M{:.2f})-{}-{} epochs -{}".format(rand_aug['N'],rand_aug['M'],model_name,epochs, run)
                     with open(res_folder+"log/%s.json" % filename, "w+") as f:
@@ -189,12 +197,13 @@ if __name__ == "__main__":
                             print('Log :\"',f.name, '\" saved !')
                         except:
                             print("Failed to save logs :",f.name)
+                            print(sys.exc_info()[1])
 
                     #plot_resV2(log, fig_name=res_folder+filename)
 
                     print('Execution Time : %.00f '%(exec_time))
                     print('-'*9)
-    #'''
+    '''
     ### HP Search ###
     '''
     from LeNet import *
@@ -221,7 +230,7 @@ if __name__ == "__main__":
 
                             t0 = time.perf_counter()
 
-                            model = getattr(models.resnet, 'resnet18')(pretrained=False)
+                            model = getattr(models.resnet, 'resnet18')(pretrained=False, num_classes=len(dl_train.dataset.classes))
                             #model = LeNet(3,10)
                             model = Higher_model(model) #run_dist_dataugV3
                             aug_model = Augmented_model(Data_augV5(TF_dict=tf_dict, N_TF=n_tf, mix_dist=dist, fixed_prob=p_setup, fixed_mag=m_setup[0], shared_mag=m_setup[1]), model).to(device)
diff --git a/higher/smart_aug/dataug.py b/higher/smart_aug/dataug.py
index 8cc91e2..e1a94b7 100755
--- a/higher/smart_aug/dataug.py
+++ b/higher/smart_aug/dataug.py
@@ -39,6 +39,7 @@ class Data_augV5(nn.Module): #Optimisation jointe (mag, proba)
             _data_augmentation (bool): Wether TF will be applied during forward pass.
             _TF_dict (dict) : A dictionnary containing the data transformations (TF) to be applied.
             _TF (list) : List of TF names.
+            _TF_ignore_mag (set): TF for which magnitude should be ignored (either it's fixed or unused).
             _nb_tf (int) : Number of TF used.
             _N_seqTF (int) : Number of TF to be applied sequentially to each inputs
             _shared_mag (bool) : Wether to share a single magnitude parameters for all TF. Beware using shared mag with basic color TF as their lowest magnitude is at PARAMETER_MAX/2.
@@ -51,7 +52,7 @@ class Data_augV5(nn.Module): #Optimisation jointe (mag, proba)
             _reg_tgt (Tensor): Target for the magnitude regularisation. Only used when _fixed_mag is set to false (ie. we learn the magnitudes).
             _reg_mask (list): Mask selecting the TF considered for the regularisation.
     """
-    def __init__(self, TF_dict=TF.TF_dict, N_TF=1, mix_dist=0.0, fixed_prob=False, fixed_mag=True, shared_mag=True):
+    def __init__(self, TF_dict, N_TF=1, mix_dist=0.0, fixed_prob=False, fixed_mag=True, shared_mag=True, TF_ignore_mag=TF.TF_ignore_mag):
         """Init Data_augv5.
 
             Args:
@@ -61,6 +62,7 @@ class Data_augV5(nn.Module): #Optimisation jointe (mag, proba)
                 fixed_prob (bool): Wether to lock the TF probabilies. (default: False)
                 fixed_mag (bool): Wether to lock the TF magnitudes. (default: True)
                 shared_mag (bool): Wether to share a single magnitude parameters for all TF. (default: True)
+                TF_ignore_mag (set): TF for which magnitude should be ignored (either it's fixed or unused).
         """
         super(Data_augV5, self).__init__()
         assert len(TF_dict)>0
@@ -71,13 +73,14 @@ class Data_augV5(nn.Module): #Optimisation jointe (mag, proba)
         #TF
         self._TF_dict = TF_dict
         self._TF= list(self._TF_dict.keys())
+        self._TF_ignore_mag=TF_ignore_mag
         self._nb_tf= len(self._TF)
         self._N_seqTF = N_TF
 
         #Mag
         self._shared_mag = shared_mag
         self._fixed_mag = fixed_mag
-        if not self._fixed_mag and len([tf for tf in self._TF if tf not in TF.TF_ignore_mag])==0:
+        if not self._fixed_mag and len([tf for tf in self._TF if tf not in self._TF_ignore_mag])==0:
             print("WARNING: Mag would be fixed as current TF doesn't allow gradient propagation:",self._TF)
             self._fixed_mag=True
 
@@ -112,7 +115,7 @@ class Data_augV5(nn.Module): #Optimisation jointe (mag, proba)
             if  self._shared_mag :
                 self._reg_tgt = torch.tensor(TF.PARAMETER_MAX, dtype=torch.float) #Encourage amplitude max
             else:
-                self._reg_mask=[self._TF.index(t) for t in self._TF if t not in TF.TF_ignore_mag]
+                self._reg_mask=[self._TF.index(t) for t in self._TF if t not in self._TF_ignore_mag]
                 self._reg_tgt=torch.full(size=(len(self._reg_mask),), fill_value=TF.PARAMETER_MAX) #Encourage amplitude max
 
     def forward(self, x):
@@ -324,6 +327,7 @@ class Data_augV7(nn.Module): #Proba sequentielles
             _data_augmentation (bool): Wether TF will be applied during forward pass.
             _TF_dict (dict) : A dictionnary containing the data transformations (TF) to be applied.
             _TF (list) : List of TF names.
+            _TF_ignore_mag (set): TF for which magnitude should be ignored (either it's fixed or unused).
             _nb_tf (int) : Number of TF used.
             _N_seqTF (int) : Number of TF to be applied sequentially to each inputs
             _shared_mag (bool) : Wether to share a single magnitude parameters for all TF. Beware using shared mag with basic color TF as their lowest magnitude is at PARAMETER_MAX/2.
@@ -336,7 +340,7 @@ class Data_augV7(nn.Module): #Proba sequentielles
             _reg_tgt (Tensor): Target for the magnitude regularisation. Only used when _fixed_mag is set to false (ie. we learn the magnitudes).
             _reg_mask (list): Mask selecting the TF considered for the regularisation.
     """
-    def __init__(self, TF_dict=TF.TF_dict, N_TF=2, mix_dist=0.0, fixed_prob=False, fixed_mag=True, shared_mag=True):
+    def __init__(self, TF_dict, N_TF=2, mix_dist=0.0, fixed_prob=False, fixed_mag=True, shared_mag=True, TF_ignore_mag=TF.TF_ignore_mag):
         """Init Data_augv7.
 
             Args:
@@ -346,6 +350,7 @@ class Data_augV7(nn.Module): #Proba sequentielles
                 fixed_prob (bool): Wether to lock the TF probabilies. (default: False)
                 fixed_mag (bool): Wether to lock the TF magnitudes. (default: True)
                 shared_mag (bool): Wether to share a single magnitude parameters for all TF. (default: True)
+                TF_ignore_mag (set): TF for which magnitude should be ignored (either it's fixed or unused).
         """
         super(Data_augV7, self).__init__()
         assert len(TF_dict)>0
@@ -359,13 +364,14 @@ class Data_augV7(nn.Module): #Proba sequentielles
         #TF
         self._TF_dict = TF_dict
         self._TF= list(self._TF_dict.keys())
+        self._TF_ignore_mag= TF_ignore_mag
         self._nb_tf= len(self._TF)
         self._N_seqTF = N_TF
 
         #Mag
         self._shared_mag = shared_mag
         self._fixed_mag = fixed_mag
-        if not self._fixed_mag and len([tf for tf in self._TF if tf not in TF.TF_ignore_mag])==0:
+        if not self._fixed_mag and len([tf for tf in self._TF if tf not in self._TF_ignore_mag])==0:
             print("WARNING: Mag would be fixed as current TF doesn't allow gradient propagation:",self._TF)
             self._fixed_mag=True
 
@@ -423,7 +429,7 @@ class Data_augV7(nn.Module): #Proba sequentielles
             if  self._shared_mag :
                 self._reg_tgt = torch.FloatTensor(TF.PARAMETER_MAX) #Encourage amplitude max
             else:
-                self._reg_mask=[idx for idx,t in enumerate(self._TF) if t not in TF.TF_ignore_mag]
+                self._reg_mask=[idx for idx,t in enumerate(self._TF) if t not in self._TF_ignore_mag]
                 self._reg_tgt=torch.full(size=(len(self._reg_mask),), fill_value=TF.PARAMETER_MAX) #Encourage amplitude max
 
     def forward(self, x):
@@ -657,7 +663,7 @@ class RandAug(nn.Module): #RandAugment = UniformFx-MagFxSh + rapide
             _fixed_mag (bool): Wether to lock the TF magnitudes. Should be True.
             _params (nn.ParameterDict): Data augmentation parameters.
     """
-    def __init__(self, TF_dict=TF.TF_dict, N_TF=1, mag=TF.PARAMETER_MAX):
+    def __init__(self, TF_dict, N_TF=1, mag=TF.PARAMETER_MAX):
         """Init RandAug.
 
             Args:
diff --git a/higher/smart_aug/test_dataug.py b/higher/smart_aug/test_dataug.py
index 24ec173..9532086 100755
--- a/higher/smart_aug/test_dataug.py
+++ b/higher/smart_aug/test_dataug.py
@@ -8,6 +8,7 @@ from dataug import *
 from train_utils import *
 
 # Use available TF (see transformations.py)
+'''
 tf_names = [
     ## Geometric TF ##
     'Identity',
@@ -57,7 +58,8 @@ tf_names = [
     #'Random',
     #'RandBlend'
 ]
-
+'''
+TF_loader=TF_loader()
 
 device = torch.device('cuda') #Select device to use
 
@@ -77,12 +79,12 @@ if __name__ == "__main__":
 
     #Task to perform
     tasks={
-        'classic',
-        #'aug_model'
+        #'classic',
+        'aug_model'
     }
     #Parameters
     n_inner_iter = 1
-    epochs = 200
+    epochs = 20
     dataug_epoch_start=0
     optim_param={
         'Meta':{
@@ -91,11 +93,11 @@ if __name__ == "__main__":
         },
         'Inner':{
             'optim': 'SGD',
-            'lr':1e-1, #1e-2/1e-1
+            'lr':1e-1, #1e-2/1e-1 (ResNet)
             'momentum':0.9, #0.9
             'decay':0.0005, #0.0005
             'nesterov':True,
-            'scheduler':'exponential', #None, 'cosine', 'multiStep', 'exponential'
+            'scheduler':'cosine', #None, 'cosine', 'multiStep', 'exponential'
         }
     }
 
@@ -137,7 +139,7 @@ if __name__ == "__main__":
         print(model_name,": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
         filename = "{}-{} epochs".format(model_name,epochs)
         #print("RandAugment-",model_name,": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
-        #filename = "RandAugment(N{}-M{:.2f})-{}-{} epochs".format(rand_aug['N'],rand_aug['M'],model_name,epochs)
+        #filename = "RandAugment(N{}-M{:.2f})-{}-{} epochs".format(rand_aug['N'],rand_aug['M'],model_name,epochs)+'-cosine'
         with open("../res/log/%s.json" % filename, "w+") as f:
             try:
                 json.dump(out, f, indent=True)
@@ -157,13 +159,23 @@ if __name__ == "__main__":
 
     #### Augmented Model ####
     if 'aug_model' in tasks:
+        tf_config='../config/base_tf_config.json'
+        tf_dict, tf_ignore_mag =TF_loader.load_TF_dict(tf_config)
+        #tf_dict = {k: TF_dict[k] for k in tf_names}
+
         torch.cuda.reset_max_memory_allocated() #reset_peak_stats
         torch.cuda.reset_max_memory_cached() #reset_peak_stats
         t0 = time.perf_counter()
 
-        tf_dict = {k: TF.TF_dict[k] for k in tf_names}
         model = Higher_model(model, model_name) #run_dist_dataugV3
-        aug_model = Augmented_model(Data_augV5(TF_dict=tf_dict, N_TF=1, mix_dist=0.5, fixed_prob=False, fixed_mag=False, shared_mag=False), model).to(device)
+        aug_model = Augmented_model(
+            Data_augV5(TF_dict=tf_dict, 
+                N_TF=3, 
+                mix_dist=0.5, 
+                fixed_prob=False, 
+                fixed_mag=False, 
+                shared_mag=False, 
+                TF_ignore_mag=tf_ignore_mag), model).to(device)
         #aug_model = Augmented_model(RandAug(TF_dict=tf_dict, N_TF=2), model).to(device)
 
         print("{} on {} for {} epochs - {} inner_it".format(str(aug_model), device_name, epochs, n_inner_iter))
@@ -175,7 +187,7 @@ if __name__ == "__main__":
              print_freq=1, 
              unsup_loss=1, 
              hp_opt=False,
-             save_sample_freq=None)
+             save_sample_freq=0)
 
         exec_time=time.perf_counter() - t0
         max_allocated = torch.cuda.max_memory_allocated()/(1024.0 * 1024.0)
@@ -188,6 +200,7 @@ if __name__ == "__main__":
             'Optimizer': optim_param, 
             "Device": device_name, 
             "Memory": [max_allocated, max_cached], 
+            "TF_config": tf_config,
             "Param_names": aug_model.TF_names(), 
             "Log": log}
         print(str(aug_model),": acc", out["Accuracy"], "in:", out["Time"][0], "+/-", out["Time"][1])
diff --git a/higher/smart_aug/train_utils.py b/higher/smart_aug/train_utils.py
index 519b838..b83863e 100755
--- a/higher/smart_aug/train_utils.py
+++ b/higher/smart_aug/train_utils.py
@@ -272,15 +272,15 @@ def run_dist_dataugV3(model, opt_param, epochs=1, inner_it=1, dataug_epoch_start
     #Scheduler
     inner_scheduler=None
     if opt_param['Inner']['scheduler']=='cosine':
-        inner_scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=epochs, eta_min=0.)
+        inner_scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(inner_opt, T_max=epochs, eta_min=0.)
     elif opt_param['Inner']['scheduler']=='multiStep':
         #Multistep milestones inspired by AutoAugment
-        inner_scheduler=torch.optim.lr_scheduler.MultiStepLR(optim, 
+        inner_scheduler=torch.optim.lr_scheduler.MultiStepLR(inner_opt, 
             milestones=[int(epochs/3), int(epochs*2/3), int(epochs*2.7/3)], 
             gamma=0.1)
     elif opt_param['Inner']['scheduler']=='exponential':
         #inner_scheduler=torch.optim.lr_scheduler.ExponentialLR(optim, gamma=0.1) #Wrong gamma
-        inner_scheduler=torch.optim.lr_scheduler.LambdaLR(optim, lambda epoch: (1 - epoch / epochs) ** 0.9)
+        inner_scheduler=torch.optim.lr_scheduler.LambdaLR(inner_opt, lambda epoch: (1 - epoch / epochs) ** 0.9)
     elif opt_param['Inner']['scheduler'] is not None:
         raise ValueError("Lr scheduler unknown : %s"%opt_param['Inner']['scheduler'])
 
diff --git a/higher/smart_aug/transformations.py b/higher/smart_aug/transformations.py
index ae2982f..dc0fc71 100755
--- a/higher/smart_aug/transformations.py
+++ b/higher/smart_aug/transformations.py
@@ -19,9 +19,9 @@ import kornia
 import random
 
 #TF that don't have use for magnitude parameter.
-TF_no_mag={'Identity', 'FlipUD', 'FlipLR', 'Random', 'RandBlend'}
+TF_no_mag={'Identity', 'FlipUD', 'FlipLR', 'Random', 'RandBlend', 'identity', 'flipUD', 'flipLR'}
 #TF which implemetation doesn't allow gradient propagaition.
-TF_no_grad={'Solarize', 'Posterize', '=Solarize', '=Posterize'}
+TF_no_grad={'Solarize', 'Posterize', '=Solarize', '=Posterize', 'posterize','solarize'}
 #TF for which magnitude should be ignored (Magnitude fixed).
 TF_ignore_mag= TF_no_mag | TF_no_grad 
 
@@ -30,6 +30,7 @@ PARAMETER_MAX = 1
 # What is the min 'level' a transform could be predicted
 PARAMETER_MIN = 0.1 
 
+'''
 # Dictionnary mapping tranformations identifiers to their function.
 # Each value of the dict should be a lambda function taking a (batch of data, magnitude of transformations) tuple as input and returns a batch of data.
 TF_dict={ #Dataugv5+
@@ -38,8 +39,8 @@ TF_dict={ #Dataugv5+
     'FlipUD' : (lambda x, mag: flipUD(x)),
     'FlipLR' : (lambda x, mag: flipLR(x)),
     'Rotate': (lambda x, mag: rotate(x, angle=rand_floats(size=x.shape[0], mag=mag, maxval=30))),
-    'TranslateX': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=x.shape[1]*0.33), zero_pos=0))),
-    'TranslateY': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=x.shape[2]*0.33), zero_pos=1))),
+    'TranslateX': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=x.shape[2]*0.33), zero_pos=0))),
+    'TranslateY': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=x.shape[3]*0.33), zero_pos=1))),
     'TranslateXabs': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=20), zero_pos=0))),
     'TranslateYabs': (lambda x, mag: translate(x, translation=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=20), zero_pos=1))),
     'ShearX': (lambda x, mag: shear(x, shear=zero_stack(rand_floats(size=(x.shape[0],), mag=mag, maxval=0.3), zero_pos=0))),
@@ -49,7 +50,7 @@ TF_dict={ #Dataugv5+
     'Contrast': (lambda x, mag: contrast(x, contrast_factor=rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.9))),
     'Color':(lambda x, mag: color(x, color_factor=rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.9))),
     'Brightness':(lambda x, mag: brightness(x, brightness_factor=rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.9))),
-    'Sharpness':(lambda x, mag: sharpeness(x, sharpness_factor=rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.9))),
+    'Sharpness':(lambda x, mag: sharpness(x, sharpness_factor=rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.9))),
     'Posterize': (lambda x, mag: posterize(x, bits=rand_floats(size=x.shape[0], mag=mag, minval=4., maxval=8.))),#Perte du gradient
     'Solarize': (lambda x, mag: solarize(x, thresholds=rand_floats(size=x.shape[0], mag=mag, minval=1/256., maxval=256/256.))), #Perte du gradient #=>Image entre [0,1]
 
@@ -57,11 +58,11 @@ TF_dict={ #Dataugv5+
     '+Contrast': (lambda x, mag: contrast(x, contrast_factor=rand_floats(size=x.shape[0], mag=mag, minval=1.0, maxval=1.9))),
     '+Color':(lambda x, mag: color(x, color_factor=rand_floats(size=x.shape[0], mag=mag, minval=1.0, maxval=1.9))),
     '+Brightness':(lambda x, mag: brightness(x, brightness_factor=rand_floats(size=x.shape[0], mag=mag, minval=1.0, maxval=1.9))),
-    '+Sharpness':(lambda x, mag: sharpeness(x, sharpness_factor=rand_floats(size=x.shape[0], mag=mag, minval=1.0, maxval=1.9))),
+    '+Sharpness':(lambda x, mag: sharpness(x, sharpness_factor=rand_floats(size=x.shape[0], mag=mag, minval=1.0, maxval=1.9))),
     '-Contrast': (lambda x, mag: contrast(x, contrast_factor=invScale_rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.0))),
     '-Color':(lambda x, mag: color(x, color_factor=invScale_rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.0))),
     '-Brightness':(lambda x, mag: brightness(x, brightness_factor=invScale_rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.0))),
-    '-Sharpness':(lambda x, mag: sharpeness(x, sharpness_factor=invScale_rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.0))),
+    '-Sharpness':(lambda x, mag: sharpness(x, sharpness_factor=invScale_rand_floats(size=x.shape[0], mag=mag, minval=0.1, maxval=1.0))),
     '=Posterize': (lambda x, mag: posterize(x, bits=invScale_rand_floats(size=x.shape[0], mag=mag, minval=4., maxval=8.))),#Perte du gradient
     '=Solarize': (lambda x, mag: solarize(x, thresholds=invScale_rand_floats(size=x.shape[0], mag=mag, minval=1/256., maxval=256/256.))), #Perte du gradient #=>Image entre [0,1]
 
@@ -86,7 +87,7 @@ TF_dict={ #Dataugv5+
     #'Auto_Contrast': (lambda mag: None), #Pas opti pour des batch (Super lent)
     #'Equalize': (lambda mag: None),
 }
-
+'''
 ## Image type cast ##
 def int_image(float_image):
     """Convert a float Tensor/Image to an int Tensor/Image.
@@ -304,7 +305,7 @@ def brightness(x, brightness_factor):
 
     return blend(torch.zeros(x.size(), device=device), x, brightness_factor).clamp(min=0.0,max=1.0) #Expect image in the range of [0, 1]
 
-def sharpeness(x, sharpness_factor):
+def sharpness(x, sharpness_factor):
     """Adjust sharpness of images.
 
     Args:
diff --git a/higher/smart_aug/utils.py b/higher/smart_aug/utils.py
index 4a4b5ae..2480268 100755
--- a/higher/smart_aug/utils.py
+++ b/higher/smart_aug/utils.py
@@ -14,6 +14,128 @@ import torch.nn.functional as F
 
 import time
 
+import transformations as TF
+class TF_loader(object):
+    """ Transformations builder.
+
+        See 'config' folder for pre-defined config files.
+
+        Attributes:
+            _filename (str): Path to config file (JSON) used.
+            _TF_dict (dict): Transformations dictionnary built from config file.
+            _TF_ignore_mag (set): Ensemble of transformations names for which magnitude should be ignored.
+            _TF_names (list): List of transformations names/keys.
+    """
+    def __init__(self):
+        """ Initialize TF_loader.
+
+        """
+        self._filename=''
+        self._TF_dict={}
+        self._TF_ignore_mag=set()
+        self._TF_names=[]
+
+    def load_TF_dict(self, filename):
+        """ Build a TF dictionnary.
+
+            Args:
+                filename (str): Path to config file (JSON) defining the transformations.
+            Returns:
+                (dict, set) : TF dicttionnary built and ensemble of TF names for which mag should be ignored.
+        """
+        self._filename=filename
+        self._TF_names=[]
+        self._TF_dict={}
+        self._TF_ignore_mag=set()
+
+        with open(filename) as json_file:
+            TF_params = json.load(json_file)
+
+        for tf in TF_params:
+            self._TF_names.append(tf['name'])
+            if tf['function'] in TF.TF_ignore_mag:
+                self._TF_ignore_mag.add(tf['name'])
+
+            if tf['function'] == 'identity':
+                self._TF_dict[tf['name']]=(lambda x, mag: x)
+
+            elif tf['function'] == 'flip':
+                #Inverser axes ?
+                if tf['param']['axis'] == 'X':
+                    self._TF_dict[tf['name']]=(lambda x, mag: TF.flipLR(x))
+                elif tf['param']['axis'] == 'Y':
+                    self._TF_dict[tf['name']]=(lambda x, mag: TF.flipUD(x))
+                else:
+                    raise Exception("Unknown TF axis : %s in %s"%(tf['function'], self._filename))
+
+            elif tf['function'] in {'translate', 'shear'}:
+                rand_fct= 'invScale_rand_floats' if tf['param']['invScale'] else 'rand_floats'
+                self._TF_dict[tf['name']]=self.build_lambda(tf['function'], rand_fct, tf['param']['min'], tf['param']['max'], tf['param']['absolute'], tf['param']['axis'])
+            
+            else:
+                rand_fct= 'invScale_rand_floats' if tf['param']['invScale'] else 'rand_floats'
+                self._TF_dict[tf['name']]=self.build_lambda(tf['function'], rand_fct, tf['param']['min'], tf['param']['max'])
+
+        return self._TF_dict, self._TF_ignore_mag
+
+    def build_lambda(self, fct_name, rand_fct_name, minval, maxval, absolute=True, axis=None):
+        """ Build a lambda function performing transformations.
+
+            Args:
+                fct_name (str): Name of the transformations to use (see transformations.py).
+                rand_fct_name (str): Name of the random mapping function to use (see transformations.py).
+                minval (float): minimum magnitude value of the TF.
+                maxval (float): maximum magnitude value of the TF.
+                absolute (bool): Wether the maxval should be relative (absolute=False) to the image size. (default: True)
+                axis (str): Axis ('X' / 'Y') of the TF, if relevant. Should be used for (flip)/translate/shear functions. (default: None)
+
+            Returns:
+                (function) transformations function : Tensor=f(Tensor, magnitude)
+        """
+        if absolute:
+            max_val_fct=(lambda x: maxval)
+        else: #Relative to img size
+            max_val_fct=(lambda x: x*maxval)
+
+        if axis is None:
+            return (lambda x, mag: 
+                        getattr(TF, fct_name)(
+                            x, 
+                            getattr(TF, rand_fct_name)(
+                                size=x.shape[0], 
+                                mag=mag, 
+                                minval=minval, 
+                                maxval=maxval)))
+        elif axis =='X':
+            return (lambda x, mag: 
+                getattr(TF, fct_name)(
+                    x, 
+                    TF.zero_stack(
+                        getattr(TF, rand_fct_name)(
+                            size=(x.shape[0],), 
+                            mag=mag, 
+                            minval=minval, 
+                            maxval=max_val_fct(x.shape[2])),
+                        zero_pos=0)))
+        elif axis == 'Y':
+            return (lambda x, mag: 
+                getattr(TF, fct_name)(
+                    x, 
+                    TF.zero_stack(
+                        getattr(TF, rand_fct_name)(
+                            size=(x.shape[0],), 
+                            mag=mag, 
+                            minval=minval, 
+                            maxval=max_val_fct(x.shape[3])),
+                        zero_pos=1)))
+        else:
+            raise Exception("Unknown TF axis : %s in %s"%(fct_name, self._filename))
+
+        def get_TF_names(self):
+            return self._TF_names
+        def get_TF_dict(self):
+            return self._TF_dict
+
 class ConfusionMatrix(object):
     """ Confusion matrix.