Override forward method in GeneralizedRCNN

I want to override forward method of GeneralizedRCNN. I want to return losses and features from ROI head during training. I have tried doing:

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(True).cuda()
model.forward = forward_modified(self,images, targets=None)

It doesn’t work quite well (throws an error from transform). Can someone share a working snippet of efficiently going about it ?

@ptrblck Can you suggest something ?

Please refer to this:

   def forward(self, images, targets=None):
            images (list[Tensor]): images to be processed
            targets (list[Dict[Tensor]]): ground-truth boxes present in the image (optional)

            result (list[BoxList] or dict[Tensor]): the output from the model.
                During training, it returns a dict[Tensor] which contains the losses.
                During testing, it returns list[BoxList] contains additional fields
                like `scores`, `labels` and `mask` (for Mask R-CNN models).

        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")
        original_image_sizes = [img.shape[-2:] for img in images]
        images, targets = self.transform(images, targets)
        features = self.backbone(images.tensors)
        if isinstance(features, torch.Tensor):
            features = OrderedDict([(0, features)])
        proposals, proposal_losses = self.rpn(images, features, targets)
        detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
        detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)

        losses = {}

        if self.training:
            return losses

        return detections