I need to implement a new architecture on the Faster-RCNN classifier on top of Faster-RCNN model implemented in Detectron2 framework. I need to implement both a new loss function on classification and a new inference logic. As reported in the official documentations (https://detectron2.readthedocs.io/en/latest/tutorials/write-models.html#:~:text=As%20an%20example%2C%20to%20use%20custom%20loss%20function%20in%20the%20box%20head%20of%20a%20Faster%20R%2DCNN%2C%20we%20can%20do%20the%20following%3A), the most flexible way should be to register a new CustomROIHead module inheriting from StandardROIHeads, and insert as box_predictor argument a custom module which inherits from FastRCNNOutputLayers. In this last, it should be possible to overwrite the function "losses" implementing the new loss as well as the other functions of interest. The problem is that I am actually not able of doing it, since I always get an error which is likely dependent on how I inherit and instantiate objects as well as how I use the super() method. I am not an expert in OOP, that's why my feeling is that the error is in there.
Here the new Module that inherits from FastRCNNOutputLayers, in which I overwrite losses() just with a new classification loss function:

class CustomRCNNOutput(FastRCNNOutputLayers):
    def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape)
    
    def losses(self, predictions, proposals):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_boxes``,
                ``gt_classes`` are expected.
        Returns:
            Dict[str, Tensor]: dict of losses
        """
        scores, proposal_deltas = predictions

        # parse classification outputs
        gt_classes = (
            cat([p.gt_classes for p in proposals], dim=0) if len(proposals) else torch.empty(0)
        )
        _log_classification_stats(scores, gt_classes)

        # parse box regression outputs
        if len(proposals):
            proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)  # Nx4
            assert not proposal_boxes.requires_grad, "Proposals should not require gradients!"
            # If "gt_boxes" does not exist, the proposals must be all negative and
            # should not be included in regression loss computation.
            # Here we just use proposal_boxes as an arbitrary placeholder because its
            # value won't be used in self.box_reg_loss().
            gt_boxes = cat(
                [(p.gt_boxes if p.has("gt_boxes") else p.proposal_boxes).tensor for p in proposals],
                dim=0,
            )
        else:
            proposal_boxes = gt_boxes = torch.empty((0, 4), device=proposal_deltas.device)

        if self.use_sigmoid_ce:
            loss_cls = self.sigmoid_cross_entropy_loss(scores, gt_classes)
        else:
            loss_cls = MY_CUSTOM_LOSS(scores, gt_classes, self.num_classes)

        losses = {
            "loss_cls": loss_cls,
            "loss_box_reg": self.box_reg_loss(
                proposal_boxes, gt_boxes, proposal_deltas, gt_classes
            ),
        }
        return {k: v * self.loss_weight.get(k, 1.0) for k, v in losses.items()}

Then, in my training script, I firstly register the new ROI Head module, I call it from in the cfg file using the reported function and after I try to run the training. (I have just reported here an extract of the training script, the other parts shouldn't be of interest).

def get_train_cfg(config_file_path, checkpoint_url, train_dataset_name, test_dataset_name, num_classes, device, output_dir):
    cfg = get_cfg() 

    cfg.merge_from_file(model_zoo.get_config_file(config_file_path))
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(checkpoint_url)
    cfg.DATASETS.TRAIN = (train_dataset_name,)
    cfg.DATASETS.TEST = (test_dataset_name,)

    cfg.DATALOADER.NUM_WORKERS = 2

    cfg.SOLVER.IMS_PER_BATCH = 2
    cfg.SOLVER.BASE_LR = 0.0002
    cfg.SOLVER.MAX_ITER = 2500
    cfg.SOLVER.STEPS = []   # for LR scheduling

    cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
    cfg.MODEL.DEVICE = device
    cfg.OUTPUT_DIR = output_dir
    cfg.MODEL.ROI_HEADS.NAME = "CustomROIHeads"

    return cfg

@ROI_HEADS_REGISTRY.register()
class CustomROIHeads(StandardROIHeads):
  def __init__(self, cfg):
    super().__init__(cfg, 
                     box_predictor=CustomRCNNOutput(cfg))

def main():
    cfg = get_train_cfg(config_file_path, checkpoint_url, train_dataset_name, test_dataset_name,
     num_classes, device, output_dir)

    with open(cfg_save_path, 'wb') as f:
        pickle.dump(cfg, f, protocol=pickle.HIGHEST_PROTOCOL)   # this will save the cfg
    
    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

    trainer = DefaultTrainer(cfg)   # call the default trainer engine and pass the custom cfg
    trainer.resume_or_load(resume=False)

    trainer.train() # to start the training

# to call the main method
if __name__ == "__main__":
    main()

Here is the error that I receive: `

Traceback (most recent call last):
  File "d:\Personale\train.py", line 107, in <module>
    main()
  File "d:\Personale\train.py", line 100, in main
    trainer = DefaultTrainer(cfg)   # call the default trainer engine and pass the custom cfg
  File "d:\personale\detectron2\detectron2\engine\defaults.py", line 376, in __init__
    model = self.build_model(cfg)
  File "d:\personale\detectron2\detectron2\engine\defaults.py", line 514, in build_model
    model = build_model(cfg)
  File "d:\personale\detectron2\detectron2\modeling\meta_arch\build.py", line 22, in build_model     
    model = META_ARCH_REGISTRY.get(meta_arch)(cfg)
  File "d:\personale\detectron2\detectron2\config\config.py", line 189, in wrapped
    explicit_args = _get_args_from_config(from_config_func, *args, **kwargs)
  File "d:\personale\detectron2\detectron2\config\config.py", line 245, in _get_args_from_config     
    ret = from_config_func(*args, **kwargs)
  File "d:\personale\detectron2\detectron2\modeling\meta_arch\rcnn.py", line 77, in from_config      
    "roi_heads": build_roi_heads(cfg, backbone.output_shape()),
  File "d:\personale\detectron2\detectron2\modeling\roi_heads\roi_heads.py", line 43, in build_roi_heads
    return ROI_HEADS_REGISTRY.get(name)(cfg, input_shape)
TypeError: CustomROIHeads.__init__() takes 2 positional arguments but 3 were given

`
Do you have any idea about how to solve my problem? It's for my Master Thesis project and currently what I did is to directly implement the new loss function on the source code and call the function in losses. It works, so it is not a problem of the algorithm, but just about how to implement it in detectron2 using a new module. This is important for me because once implementing also the new inference mode, it will be not flexible to implement it directly in the source code. Also for sharing, it would be impossible. I need to find the correct way to introduce a new module.
Thank you in advance!