What is the best way of evaluating a model exported with torch.jit.trace?

[mikel-brostrom]

I am exporting the model by:

ts = torch.jit.trace(model, (im, im), strict=False)
optimize_for_mobile(ts)._save_for_lite_interpreter('/path/to/model.torchscript')

Then I load it in a LightningModule with exactly the same validation_step as the original model:

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()

        # init a pretrained resnet
        self.model = torch.jit.load('/path/to/model.torchscript')

        self.f1 = F1Score(ignore_index=0, num_classes=2, average='macro', mdmc_average='samplewise')
        self.recall = Recall(ignore_index=0, num_classes=2, average='macro', mdmc_average='samplewise')
        self.precision_ = Precision(ignore_index=0, num_classes=2, average='macro', mdmc_average='samplewise')  # self.precision exists in 

    def forward(self, t0, t1):
        x = self.model(t0, t1)
        return x

    def validation_step(self, batch, batch_idx):

        t0_orig, t1_orig, t0, t1, mask_gt = batch
        mask_pred = self.forward(t0, t1)

        mask_pred = torch.where(mask_pred > 0.5, 1, 0)

        f1_score = self.f1(mask_pred, mask_gt)
        precision = self.precision_(mask_pred, mask_gt)
        recall = self.recall(mask_pred, mask_gt)

        stat_scores_ = stat_scores(mask_pred, mask_gt, num_classes=2, ignore_index=0, reduce='macro', mdmc_reduce='samplewise')
        # extract stat_scores for CHANGE class, ignore NO CHANGE
        stat_scores_ = stat_scores_[:,1,:]
        # reduce over batch
        stat_scores_ = torch.sum(stat_scores_, dim=0)
        tp = stat_scores_[0].to(torch.float)
        fp = stat_scores_[1].to(torch.float)
        tn = stat_scores_[2].to(torch.float)
        fn = stat_scores_[3].to(torch.float)

        # gpu tensor --> cpu numpy array
        mask_pred = mask_pred.detach().cpu().numpy()
        mask_pred = np.where(mask_pred > 0.5, 255, 0)

        metrics = {
            'val_f1_score': f1_score,
            'val_precision': precision,
            'val_recall': recall,
            'val_TP': tp,
            'val_FP': fp,
            'val_TN': tn,
            'val_FN': fn,
        }

        self.log_dict(metrics, on_epoch=True)

and finally I setup my dm and trainer:

model = Model()
dm = DataModule(config)
trainer = pl.Trainer(
        accelerator='cpu',
    )
trainer.validate(model, datamodule=dm)

Is there a better way f doing this?