remove manual no_grad and train mode

airctic · May 16, 2022 · 13bb3e7 · 13bb3e7
1 parent 3dbbb44
commit 13bb3e7
Showing 1 changed file with 9 additions and 10 deletions.
diff --git 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)