Support training with the vanilla module

test/test_engine.py

@@ -95,7 +95,6 @@ def test_train_with_vanilla_module(self):
         self.assertIsInstance(out["bbox_regression"], torch.Tensor)
         self.assertIsInstance(out["objectness"], torch.Tensor)
 
-    @unittest.skip("Just ignore this.")
     def test_train_one_step(self):
         # Load model
         model = yolov5s()

yolort/models/yolo_module.py

@@ -77,9 +77,31 @@ def forward(
         # Transform the input
         samples, targets = self.transform(inputs, targets)
         # Compute the detections
-        detections = self.model(samples.tensors, targets=targets)
-        # Rescale coordinate
-        detections = self.transform.postprocess(detections, samples.image_sizes, original_image_sizes)
+        outputs = self.model(samples.tensors, targets=targets)
+
+        losses = {}
+        detections: List[Dict[str, Tensor]] = []
+
+        if self.training:
+            # compute the losses
+            losses = outputs
+        else:
+            # Rescale coordinate
+            detections = self.transform.postprocess(outputs, samples.image_sizes, original_image_sizes)
+
+        if torch.jit.is_scripting():
+            return losses, detections
+        else:
+            return self.eager_outputs(losses, detections)
+
+    @torch.jit.unused
+    def eager_outputs(
+        self,
+        losses: Dict[str, Tensor],
+        detections: List[Dict[str, Tensor]],
+    ) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]:
+        if self.training:
+            return losses
 
         return detections