fix loss logging issue

SeungyounShin · May 23, 2023 · 4a85b17 · 4a85b17
1 parent a993949
commit 4a85b17
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/lavis/tasks/base_task.py b/lavis/tasks/base_task.py
@@ -61,9 +61,12 @@ def build_datasets(self, cfg):
         return datasets
 
     def train_step(self, model, samples):
-        loss_dict = model(samples)
-        loss = loss_dict["loss"]
-        return loss, loss_dict
+        output = model(samples)
+        loss_dict = {}
+        for k,v in output.items():
+            if "loss" in k:
+                loss_dict[k] = v
+        return output["loss"], loss_dict
 
     def valid_step(self, model, samples):
         raise NotImplementedError