Apply suggestions from code review

Co-authored-by: Nicki Skafte <skaftenicki@gmail.com>
Lightning-Universe · Sep 24, 2021 · 857a19f · 857a19f
1 parent 10a4fde
commit 857a19f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/pl_bolts/models/detection/faster_rcnn/faster_rcnn_module.py b/pl_bolts/models/detection/faster_rcnn/faster_rcnn_module.py
@@ -99,6 +99,7 @@ def training_step(self, batch, batch_idx):
  loss_dict = self.model(images, targets)
  loss = sum(loss for loss in loss_dict.values())
  self.log("loss", loss, prog_bar=True)
+ return loss
 
  def validation_step(self, batch, batch_idx):
  images, targets = batch

diff --git a/pl_bolts/models/detection/retinanet/retainanet_module.py b/pl_bolts/models/detection/retinanet/retainanet_module.py
@@ -92,13 +92,15 @@ def training_step(self, batch, batch_idx):
  loss_dict = self.model(images, targets)
  loss = sum(loss for loss in loss_dict.values())
  self.log("loss", loss, prog_bar=True)
+ return loss
 
  def validation_step(self, batch, batch_idx):
  images, targets = batch
  # fasterrcnn takes only images for eval() mode
  outs = self.model(images)
  iou = torch.stack([_evaluate_iou(t, o) for t, o in zip(targets, outs)]).mean()
  self.log("val_iou", iou, prog_bar=True)
+ return {"val_iou": iou}
 
  def validation_epoch_end(self, outs):
  avg_iou = torch.stack([o["val_iou"] for o in outs]).mean()