diff --git a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.024e-7_fpn.yaml b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.024e-7_fpn.yaml
index a6d33a5..0587784 100644
--- a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.024e-7_fpn.yaml
+++ b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.024e-7_fpn.yaml
@@ -116,11 +116,11 @@ train:
       milestones: [5]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    func2extract_org_loss: 'extract_org_loss_dict'
-    org_term:
-      weight: 1.0
-    sub_terms:
+    key: 'WeightedSumLoss'
+    func2extract_model_loss: 'extract_model_loss_dict'
+    kwargs:
+      model_term:
+        weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.28e-8_fpn.yaml b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.28e-8_fpn.yaml
index f50cafd..feeaac2 100644
--- a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.28e-8_fpn.yaml
+++ b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta1.28e-8_fpn.yaml
@@ -116,11 +116,11 @@ train:
       milestones: [5]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    func2extract_org_loss: 'extract_org_loss_dict'
-    org_term:
-      weight: 1.0
-    sub_terms:
+    key: 'WeightedSumLoss'
+    func2extract_model_loss: 'extract_model_loss_dict'
+    kwargs:
+      model_term:
+        weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta2.048e-7_fpn.yaml b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta2.048e-7_fpn.yaml
index 45661ca..f517c7e 100644
--- a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta2.048e-7_fpn.yaml
+++ b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta2.048e-7_fpn.yaml
@@ -116,11 +116,11 @@ train:
       milestones: [5]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    func2extract_org_loss: 'extract_org_loss_dict'
-    org_term:
-      weight: 1.0
-    sub_terms:
+    key: 'WeightedSumLoss'
+    func2extract_model_loss: 'extract_model_loss_dict'
+    kwargs:
+      model_term:
+        weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta3.2768e-6_fpn.yaml b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta3.2768e-6_fpn.yaml
index 069dbc9..005cbea 100644
--- a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta3.2768e-6_fpn.yaml
+++ b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta3.2768e-6_fpn.yaml
@@ -116,11 +116,11 @@ train:
       milestones: [5]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    func2extract_org_loss: 'extract_org_loss_dict'
-    org_term:
-      weight: 1.0
-    sub_terms:
+    key: 'WeightedSumLoss'
+    func2extract_model_loss: 'extract_model_loss_dict'
+    kwargs:
+      model_term:
+        weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta8.192e-7_fpn.yaml b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta8.192e-7_fpn.yaml
index 50816a9..899d188 100644
--- a/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta8.192e-7_fpn.yaml
+++ b/configs/coco2017/supervised_compression/end-to-end/faster_rcnn_splittable_resnet50-fp-beta8.192e-7_fpn.yaml
@@ -116,11 +116,11 @@ train:
       milestones: [5]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    func2extract_org_loss: 'extract_org_loss_dict'
-    org_term:
-      weight: 1.0
-    sub_terms:
+    key: 'WeightedSumLoss'
+    func2extract_model_loss: 'extract_model_loss_dict'
+    kwargs:
+      model_term:
+        weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.08_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.08_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 61441d3..f7620bf 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.08_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.08_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.16_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.16_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 358f9bc..954e7cc 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.16_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.16_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.32_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.32_fpn_from_faster_rcnn_resnet50_fpn.yaml
index fe81510..833a5d4 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.32_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.32_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.64_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.64_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 480406b..a6ac5a3 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.64_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta0.64_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta1.28_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta1.28_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 2b7968a..5013af7 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta1.28_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta1.28_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta2.56_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta2.56_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 3c50a33..2297e37 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta2.56_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta2.56_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta5.12_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta5.12_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 90447e9..434e5e7 100644
--- a/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta5.12_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/entropic_student/faster_rcnn_splittable_resnet50-fp-beta5.12_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -140,70 +140,77 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      org_term:
-        weight: 0.0
-      sub_terms:
-        layer1:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.bottleneck_layer'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer1'
-              io: 'output'
-          weight: 1.0
-        layer2:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer2'
-              io: 'output'
-          weight: 1.0
-        layer3:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer3'
-              io: 'output'
-          weight: 1.0
-        layer4:
-          criterion:
-            key: 'MSELoss'
-            kwargs:
-              reduction: 'sum'
-          kwargs:
-            input:
-              is_from_teacher: False
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-            target:
-              is_from_teacher: True
-              module_path: 'seq.backbone.body.layer4'
-              io: 'output'
-          weight: 1.0
+      key: 'WeightedSumLoss'
+      kwargs:
+        sub_terms:
+          layer1:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.bottleneck_layer'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer1'
+                  io: 'output'
+            weight: 1.0
+          layer2:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer2'
+                  io: 'output'
+            weight: 1.0
+          layer3:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer3'
+                  io: 'output'
+            weight: 1.0
+          layer4:
+            criterion:
+              key: 'MSELoss'
+              kwargs:
+                reduction: 'sum'
+            criterion_wrapper:
+              key: 'SimpleLossWrapper'
+              kwargs:
+                input:
+                  is_from_teacher: False
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+                target:
+                  is_from_teacher: True
+                  module_path: 'seq.backbone.body.layer4'
+                  io: 'output'
+            weight: 1.0
   stage2:
     num_epochs: 5
     train_data_loader:
@@ -254,11 +261,11 @@ train:
         milestones: [3]
         gamma: 0.1
     criterion:
-      key: 'GeneralizedCustomLoss'
-      func2extract_org_loss: 'extract_org_loss_dict'
-      org_term:
-        weight: 1.0
-      sub_terms:
+      key: 'WeightedSumLoss'
+      func2extract_model_loss: 'extract_model_loss_dict'
+      kwargs:
+        model_term:
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq12ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq12ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 4e87dd2..8e2e9a5 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq12ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq12ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq1ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq1ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 3564f82..a194ba2 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq1ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq1ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq2ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq2ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 994e401..1b48d5b 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq2ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq2ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq3ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq3ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 557236a..9cf9a97 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq3ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq3ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq6ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq6ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 2a1b679..d2d1638 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq6ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq6ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader:
diff --git a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq9ch_fpn_from_faster_rcnn_resnet50_fpn.yaml b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq9ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
index 71588c8..744b5e9 100644
--- a/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq9ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
+++ b/configs/coco2017/supervised_compression/ghnd-bq/faster_rcnn_resnet50-bq9ch_fpn_from_faster_rcnn_resnet50_fpn.yaml
@@ -159,70 +159,77 @@ train:
       milestones: [5, 15]
       gamma: 0.1
   criterion:
-    key: 'GeneralizedCustomLoss'
-    org_term:
-      weight: 0.0
-    sub_terms:
-      layer1:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.bottleneck_layer'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer1'
-            io: 'output'
-        weight: 1.0
-      layer2:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer2'
-            io: 'output'
-        weight: 1.0
-      layer3:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer3'
-            io: 'output'
-        weight: 1.0
-      layer4:
-        criterion:
-          key: 'MSELoss'
-          kwargs:
-            reduction: 'sum'
-        kwargs:
-          input:
-            is_from_teacher: False
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-          target:
-            is_from_teacher: True
-            module_path: 'seq.backbone.body.layer4'
-            io: 'output'
-        weight: 1.0
+    key: 'WeightedSumLoss'
+    kwargs:
+      sub_terms:
+        layer1:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.bottleneck_layer'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer1'
+                io: 'output'
+          weight: 1.0
+        layer2:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer2'
+                io: 'output'
+          weight: 1.0
+        layer3:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer3'
+                io: 'output'
+          weight: 1.0
+        layer4:
+          criterion:
+            key: 'MSELoss'
+            kwargs:
+              reduction: 'sum'
+          criterion_wrapper:
+            key: 'SimpleLossWrapper'
+            kwargs:
+              input:
+                is_from_teacher: False
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+              target:
+                is_from_teacher: True
+                module_path: 'seq.backbone.body.layer4'
+                io: 'output'
+          weight: 1.0
 
 test:
   test_data_loader: