Regarding "forward" for custom models

Hellow there.

I’m currently trying to build a custom object detection model based on Faster R-CNN.
The structure of the custom CNN model is such that the final output is the class of the object, the boundary box, and the attributes of the object .
So, I wrote the following code.

class CustomPredictor(nn.Module):
    def __init__(self,in_channels,num_classes):
        self.cls_score = nn.Linear(in_channels, num_classes)
        self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
        self.additional_layer = nn.Linear(in_channels,1) #this is the additional layer
    def forward(self,x):
        if x.dim() == 4:
            assert list(x.shape[2:]) == [1, 1]
        x = x.flatten(start_dim=1)
        pred_leaf_age = self.additional_layer(x)
        scores = self.cls_score(x)
        bbox_deltas = self.bbox_pred(x)
        return scores, bbox_deltas, pred_leaf_age

class CustomModel(nn.Module):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
        self.in_features = self.model.roi_heads.box_head.fc7.out_features
        self.model.roi_heads.box_predictor = CustomPredictor(self.in_features,6)
    def forward(self, x):

The problem here is how to define the foward of the CustomModel class.
Could you give me any advice?

Best regards.

I may have resolved this probrem in myself .
Thank you for watching.