Quantize Aware Training - FasterRCNN

[suchig]

I tried to quantize a FasterRCNN based model (from pytorch torchvision). My Lightning module is shown way below.

Problem:
When I try Trainer.fit(), I get the error that "one of the variables needed for gradient computation has been modified by an inplace operation".

What I tried:

1. I was able to nail down that this is specifically happening because of losses. There are 4 losses that make up loss for pytorch based fatserrcnn - detector losses (loss_classification and loss_box_reg) and proposal based losses (loss_objectness and loss_rpn_box_reg)
2. Upon debugging further, I found that this is specifically caused by proposal based losses (loss_objectness and loss_rpn_box_reg). Question 1: Why are these losses throwing the above error? I saw the implementation and could not detect any inplace operations and they look similar to detector based losses
3. Instead of returning sum of losses in training step, I returned just the sum of detector losses (loss_classification and loss_box_reg). This helped me continue with Training. But the loss started increasing and exploding and the training crashed after a point.

Ask:
How to resolve this error?

Quant Model:

class FasterRCNN_Quant(pl.LightningModule):
def init(self, model: torch.nn.Module ):
super().init()
self.model = model

def forward(self, x):
    return self.model(x)

def training_step(self, batch, batch_idx):
    x, y = batch 
    loss_dict = self.model(x, y)
    losses = sum(loss for loss in loss_dict.values())
    
    return losses
def configure_optimizers(self):
    optimizer = torch.optim.SGD(self.model.parameters(),
                                lr=0.005,
                                momentum=0.9,
                                weight_decay=0.005)
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                           step_size=3,
                                           gamma=0)
    return [optimizer], [lr_scheduler]