fix loss logging issue

salesforce · 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