adjustment fot CCS

models/cnn_baseline.py

@@ -14,10 +14,12 @@ class CNNBaseline(CombinedModelMixins,
                   LightningBaseModule
                   ):
 
-    def __init__(self, in_shape, n_classes, weight_init, activation, use_bias, use_norm, dropout, lat_dim, features,
+    def __init__(self, in_shape, n_classes, weight_init, activation,
+                 use_bias, use_norm, dropout, lat_dim, features,
                  filters,
                  lr, weight_decay, sto_weight_avg, lr_warm_restart_epochs, opt_reset_interval,
-                 loss, scheduler):
+                 loss, scheduler, lr_scheduler_parameter
+                 ):
 
         # TODO: Move this to parent class, or make it much easieer to access....
         a = dict(locals())