I applied Quantisation aware training using PyTorch lightning on one of the architectures for faster inference, The model has been trained successfully but I am facing model loading issues during inference. I’ve come across a few forums with this same issue but couldn’t find a satisfactory method that can resolve my issue. Any help would be highly appreciated, Thanks!
Below is the attached code (I have trained the QAT model using PyTorch lightning but the issue arises when I try to load it)
Training Code
import torch
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks import QuantizationAwareTraining


class Module(pl.LightningModule):

    def __init__(self, model):
        super().__init__()
        self.save_hyperparameters()

        self.model = Model(3).to(device)

        self.lr = lr
  

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

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

    def training_step(self, batch, batch_idx):
        lr, hr = batch
        sr = self(lr)
        loss = F.mse_loss(sr, hr, reduction="mean")

        return loss

    def validation_step(self, batch, batch_idx):
        lr, hr = batch
        sr = self(lr)
        loss = F.mse_loss(sr, hr, reduction="mean")

        return loss

    def test_step(self, batch, batch_idx):
        
        lr, hr = batch
        sr = self(lr)
        loss = F.mse_loss(sr, hr, reduction="mean")

        return loss


if __name__ == '__main__':


    scale_factor = 3
    batch_size = 24
            
    epochs = 1
    lr = 1e-5
        
    input_image_path = '....'  
    target_image_path = '....'

    val_input_path = '....'
    val_target_path = '...'
    
    prev_ckpt_path = '...'
    device = 'cpu' # Device kept as CPU for Quantisation Aware Training, as it doesnt support GPU
    
    # Define model
    model = SRModel(scale_factor).to('cpu')
    
    module = Module(model).load_from_checkpoint(prev_ckpt_path)

    # Setup dataloaders


    train_dataset = CustomDataset(input_image_path, target_image_path)
    
    training_dataloader = DataLoader(
                dataset=train_dataset, num_workers=4, batch_size=batch_size, shuffle=True)

    val_dataset = CustomDataset(val_input_path, val_target_path)
    
    val_dataloader = DataLoader(
        dataset=val_dataset, num_workers=4, batch_size=batch_size, shuffle=False)
    

    checkpoint_callback = ModelCheckpoint(monitor='val_loss')

    trainer = pl.Trainer(max_epochs=epochs, gpus=0,  auto_lr_find=True,                         
                        logger= wandb_logger, progress_bar_refresh_rate = 3,
                        callbacks=[QuantizationAwareTraining(observer_type='histogram', input_compatible=True), checkpoint_callback])


    trainer.fit(
        module,
        training_dataloader,
        val_dataloader
    )
    

    trainer.save_checkpoint("Quantised.pth")
    trainer.save_checkpoint("Quantised.ckpt")

And this is the inference code 

Inference


class Module(pl.LightningModule):

    def __init__(self, model):
        super().__init__()
        self.save_hyperparameters()

        self.model = SRModel(3).to(device)

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


prev_ckpt_path = '.....ckpt'

device = 'cpu'

# Define model
model = SRModel(3).to(device)

module = Module(model).load_from_checkpoint(prev_ckpt_path, strict=False)


RunTime Error: an exception occurred : (‘Copying from quantized Tensor to non-quantized Tensor is not allowed, please use dequantize to get a float Tensor from a quantized Tensor’,).