Different resutls for trained model and model loaded from checkpoint

[d3mage]

Hi. I am currently working with brain tumor segmentation with a custom dataset. During the training, my model has been outputting good results for training, testing, and validation data. However, when I load the model with the checkpoint, the same testing data from the training step now produces bad results.
I have tried to resolve the issue by experimenting with the backbone, encoders, optimizers, hyperparameters, etc. Nothing worked and, lastly, I have conducted an experiment where I trained a small model and forwarded dummy data expecting to get the same output.

Training loop

BATCH_SIZE = 4
EPOCHS = 5


for i, (train_index, test_index) in enumerate(kf.split(patients)):
    train_patients, test_patients = (
        [patients[index] for index in train_index],
        [patients[index] for index in test_index],
    )

    total_volumes = []

    train_ds = BrainDataset(train_patients)

    train_dataloader = (
        torch.utils.data.DataLoader(
            train_ds,
            batch_size=BATCH_SIZE,
            shuffle=True,
        ),
    )

    model = Model("deeplabv3plus", "resnet50", in_channels=1, out_classes=1)

    trainer = L.Trainer(max_epochs=EPOCHS, enable_checkpointing=False)

    trainer.fit(
        model,
        train_dataloaders=train_dataloader,
        val_dataloaders=valid_dataloader,
    )

    trainer.save_checkpoint(f"./weights/temp/5.ckpt")

    break

	result = model(torch.ones([4, 1, 256, 256]))

	print(result)

tensor([[[[-1.1145, -1.0780, -1.0416,  ..., -0.7698, -0.7786, -0.7873],

Loading model

model = Model.load_from_checkpoint("weights/temp/5.ckpt")

result = model(torch.ones([4, 1, 256, 256]))

print(result)

tensor([[[[-1.4208, -1.1938, -0.9668,  ..., -0.5022, -0.4182, -0.3343],

Model

import torch
import torch.nn as nn
import torch.optim as optim
import segmentation_models_pytorch as smp
import lightning as L


class Model(L.LightningModule):
    def __init__(self, arch, encoder_name, in_channels, out_classes, **kwargs):
        super().__init__()
        self.save_hyperparameters()

        self.model = smp.create_model(
            arch,
            encoder_name=encoder_name,
            in_channels=in_channels,
            classes=out_classes,
            **kwargs,
        )

        params = smp.encoders.get_preprocessing_params(encoder_name)

        mean = torch.tensor(params["mean"][0], dtype=torch.float32)
        std = torch.tensor(params["std"][0], dtype=torch.float32)

        self.register_buffer("mean", mean.view(1, in_channels, 1, 1))
        self.register_buffer("std", std.view(1, in_channels, 1, 1))

        self.dice_loss = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True)
        self.ce_loss = nn.BCEWithLogitsLoss()
        self.loss_weight_dice = 0.8
        self.loss_weight_ce = 0.2

        self.outputs = []

    def forward(self, image):
        image = (image - self.mean) / self.std

        mask = self.model(image)
        return mask

    def loss_fn(self, logits_mask, mask):
        loss = self.loss_weight_dice * self.dice_loss(
            logits_mask, mask
        ) + self.loss_weight_ce * self.ce_loss(logits_mask, mask)
        return loss

    def training_step(self, batch, batch_idx):
        return self.shared_step(batch, "train")

    def validation_step(self, batch, batch_idx):
        return self.shared_step(batch, "valid")

    def test_step(self, batch, batch_idx):
        return self.shared_step(batch, "test")

    def shared_step(self, batch, stage):
        if isinstance(batch, tuple):
            batch = batch[0]

        assert isinstance(batch, dict), f"Batch is not a dictionary but {type(batch)}"

        image, mask = batch.get("image"), batch.get("mask")

        assert image.ndim == 4 and mask.ndim == 4
        h, w = image.shape[2:]
        assert h % 32 == 0 and w % 32 == 0
        assert mask.max() <= 1.0 and mask.min() >= 0

        logits_mask = self.forward(image)

        loss = self.loss_fn(logits_mask, mask)

        # Metrics calculation
        prob_mask = logits_mask.sigmoid()
        pred_mask = (prob_mask > 0.5).float()
        tp, fp, fn, tn = smp.metrics.get_stats(
            pred_mask.long(), mask.long(), mode="binary"
        )

        output = {
            "loss": loss,
            "tp": tp,
            "fp": fp,
            "fn": fn,
            "tn": tn,
        }

        self.outputs.append(output)

        if stage == "train":
            self.on_train_epoch_end()
        elif stage == "valid":
            self.on_validation_epoch_end()
        elif stage == "test":
            self.on_test_epoch_end()
        return output

    def on_shared_epoch_end(self, stage):
        tp = torch.cat([x["tp"] for x in self.outputs])
        fp = torch.cat([x["fp"] for x in self.outputs])
        fn = torch.cat([x["fn"] for x in self.outputs])
        tn = torch.cat([x["tn"] for x in self.outputs])

        per_image_iou = smp.metrics.iou_score(
            tp, fp, fn, tn, reduction="micro-imagewise"
        )
        dataset_iou = smp.metrics.iou_score(tp, fp, fn, tn, reduction="micro")

        metrics = {
            f"{stage}_per_image_iou": per_image_iou,
            f"{stage}_dataset_iou": dataset_iou,
        }

        self.log_dict(metrics, prog_bar=True)

    def on_train_epoch_end(self):
        self.on_shared_epoch_end("train")

    def on_validation_epoch_end(self):
        self.on_shared_epoch_end("valid")

    def on_test_epoch_end(self):
        self.on_shared_epoch_end("test")

    def configure_optimizers(self):
        # optimizer = torch.optim.SGD(self.parameters(), lr=0.1, momentum=0.7, weight_decay=1e-6)
        optimizer = torch.optim.Adam(self.parameters(), lr=0.1e-4)
        lr_scheduler = {
            "scheduler": optim.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                mode="max",
                factor=0.1,
                patience=3,
                threshold=0.0001,
                threshold_mode="rel",
                verbose=True,
                cooldown=5,
            ),
            "monitor": "train_per_image_iou",
            "frequency": 1,
        }
        return {"optimizer": optimizer}

Why do I get different results? How do I make the results consistent? Thank you in advance.