Additional argument in Faster-RCNN forwardpass

Hello everyone,

I am currently trying to boost the classification accuracy of Faster-RCNN by using and auxiliary classification network for feature enrichment. To utilize the new inputs I need to add an additional argument in my box_predictor where x_cls is the last flatten layer of my auxiliary classification network.

class myRichDetector(nn.Module):

    def __init__(self, in_channels, cls_flatten, num_classes):
        super(myRichDetector, self).__init__()
        # size of cls_flatten is equal to the size of the last flatten layer of the aux network
        enriched_layer = in_channels + cls_flatten
        self.pred_tor = FastRCNNPredictor(in_channels, num_classes)
        self.pred_tor.cls_score = nn.Linear(enriched_layer, num_classes)
        self.pred_tor.bbox_pred = nn.Linear(in_channels, num_classes * 4)

    def forward(self, x, x_cls):
        # enrich the output of the F-RCNN with the aux output
        x_enr = torch.torch.cat((x, x_cls), 1)
        if x.dim() == 4:
            assert list(x.shape[2:]) == [1, 1]
        x = x.flatten(start_dim=1)
        scores = self.pred_tor.cls_score(x_enr)
        bbox_deltas = self.pred_tor.bbox_pred(x)

        return scores, bbox_deltas

class Model_detection(nn.Module):
    def __init__(self, num_classes= None, in_features_cls=None, feature_extract=True, use_pretrained=True):
        super(Model_detection, self).__init__()


        self.model = models.detection.fasterrcnn_resnet50_fpn(pretrained=use_pretrained)
        in_features = self.model.roi_heads.box_predictor.cls_score.in_features

        # slice the model to extract the roi_heads and remove the box_predictor
        self.model_my_roi_heads = nn.Sequential(*list(self.model.roi_heads.children())[:-1])
        # remove the roi_head from the model
        self.model = nn.Sequential(*list(self.model.children())[:-1])
        # define the custom box_predictor
        self.model_rich_head = myRichDetector(in_features, in_features_cls, num_classes)

        self.set_parameter_requires_grad(self.model, feature_extract)
        self.set_parameter_requires_grad(self.model_my_roi_heads, feature_extract)
        self.set_parameter_requires_grad(self.model_rich_head, feature_extract)


    def set_parameter_requires_grad(self, model, feature_extracting):
        if feature_extracting:
            for param in model.parameters():
                param.requires_grad = False

    def forward(self, x, target, cls_latent):
        
        # normal F-RCNN forward pass
        x = self.model(x, target)
        x = self.model_my_roi_heads(x)
        
        # introduce the new features for the prediction
        x = self.model_rich_head(x, cls_latent)

        return x

Now x = self.model(x, target) gives me following error:

TypeError: forward() takes 2 positional arguments but 3 were given

Can anybody help me with this issue?
Thank you very much in advance