🛠 Fix the loss-val (#425) in STFPM

openvinotoolkit · Jul 11, 2022 · 00492cf · 00492cf
1 parent 7e34bb1
commit 00492cf
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/anomalib/models/stfpm/lightning_model.py b/anomalib/models/stfpm/lightning_model.py
@@ -56,7 +56,6 @@ def __init__(
             layers=layers,
         )
         self.loss = STFPMLoss()
-        self.loss_val = 0
 
     def training_step(self, batch, _):  # pylint: disable=arguments-differ
         """Training Step of STFPM.
@@ -72,8 +71,7 @@ def training_step(self, batch, _):  # pylint: disable=arguments-differ
         """
         self.model.teacher_model.eval()
         teacher_features, student_features = self.model.forward(batch["image"])
-        loss = self.loss_val + self.loss(teacher_features, student_features)
-        self.loss_val = 0
+        loss = self.loss(teacher_features, student_features)
         return {"loss": loss}
 
     def validation_step(self, batch, _):  # pylint: disable=arguments-differ