From 13bb3e72279cdc5aa22d047c7e33d2c05ec31861 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Mon, 16 May 2022 17:13:57 +0530 Subject: [PATCH] remove manual no_grad and train mode --- .../yolov5/lightning/model_adapter.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/icevision/models/ultralytics/yolov5/lightning/model_adapter.py b/icevision/models/ultralytics/yolov5/lightning/model_adapter.py index f967f097e..5a9ec93d6 100644 --- a/icevision/models/ultralytics/yolov5/lightning/model_adapter.py +++ b/icevision/models/ultralytics/yolov5/lightning/model_adapter.py @@ -42,16 +42,15 @@ def training_step(self, batch, batch_idx): def validation_step(self, batch, batch_idx): (xb, yb), records = batch - with torch.no_grad(): - inference_out, training_out = self(xb) - preds = yolov5.convert_raw_predictions( - batch=xb, - raw_preds=inference_out, - records=records, - detection_threshold=0.001, - nms_iou_threshold=0.6, - ) - loss = self.compute_loss(training_out, yb)[0] + inference_out, training_out = self(xb) + preds = yolov5.convert_raw_predictions( + batch=xb, + raw_preds=inference_out, + records=records, + detection_threshold=0.001, + nms_iou_threshold=0.6, + ) + loss = self.compute_loss(training_out, yb)[0] self.accumulate_metrics(preds)