Merge pull request #2249 from f4str/pytorch-frcnn-loss-fix

Fix PyTorch Object Detection Estimators Missing Gradients Bug

Author: beat-buesser

art/estimators/object_detection/pytorch_object_detector.py

@@ -270,6 +270,12 @@ def _get_losses(
         x_preprocessed = x_preprocessed.to(self.device)
         y_preprocessed = [{k: v.to(self.device) for k, v in y_i.items()} for y_i in y_preprocessed]
 
+        # Set gradients again after inputs are moved to another device
+        if x_preprocessed.is_leaf:
+            x_preprocessed.requires_grad = True
+        else:
+            x_preprocessed.retain_grad()
+
         loss_components = self._model(x_preprocessed, y_preprocessed)
 
         return loss_components, x_preprocessed

art/estimators/object_detection/pytorch_yolo.py

@@ -358,6 +358,12 @@ def _get_losses(
         x_preprocessed = x_preprocessed.to(self.device)
         y_preprocessed_yolo = y_preprocessed_yolo.to(self.device)
 
+        # Set gradients again after inputs are moved to another device
+        if x_preprocessed.is_leaf:
+            x_preprocessed.requires_grad = True
+        else:
+            x_preprocessed.retain_grad()
+
         # Calculate loss components
         loss_components = self._model(x_preprocessed, y_preprocessed_yolo)