Getting all Nan values in prediction when testing model loaded from checkpoint

[valecorbetta]

Hi,
I have trained a UNetr model for semantic segmentation of 3D CT images. I wanted to test the my test set, however when run the model on it the predictions only contain Nan values.

The model is defined as follows:

class Model(pl.LightningModule):
    def __init__(self, net, criterion, learning_rate, optimizer_class):
        super().__init__()
        self.lr = learning_rate
        self.net = net
        self.criterion = criterion
        self.optimizer_class = optimizer_class
        self.dice = torchmetrics.functional.dice_score

    def configure_optimizers(self):
        optimizer = self.optimizer_class(self.parameters(), lr=self.lr)
        return optimizer

    def prepare_batch(self, batch):
        return batch['image'][tio.DATA], batch['label'][tio.DATA]

    def infer_batch(self, batch):
        x, y = self.prepare_batch(batch)
        y_hat = self.net(x)
        return y_hat, y

    def training_step(self, batch, batch_idx):
        y_hat, y = self.infer_batch(batch)
        loss = self.criterion(y_hat, y)
        preds = self(y_hat)
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_dice_step', self.dice(preds, y), prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        y_hat, y = self.infer_batch(batch)
        loss = self.criterion(y_hat, y)
        preds = self(y_hat)
        self.log('val_loss', loss)
        self.log('val_dice_step', self.dice(preds,y), prog_bar=True)
        return loss

And the test script is as follows:

device = "gpu"
ROOT = Path.cwd()

unetr = monai.networks.nets.UNETR(
        in_channels=1,
        out_channels=1,
        img_size=(32, 32, 32),
        feature_size=32,
        hidden_size=768,
        mlp_dim=3072,
        num_heads=12,
        pos_embed="conv",
        norm_name="instance",
        res_block=True,
        dropout_rate=0.0,
    )

saved_chk = ROOT / "best_metric_model.ckpt"
model_dict = til.Model.load_from_checkpoint(saved_chk, net =unetr, criterion=monai.losses.DiceCELoss(sigmoid=True), learning_rate=0.00001, optimizer_class=torch.optim.AdamW)
model_dict.eval()
model_dict.to(device)
validation_set = data.test_set
subject = random.choice(validation_set)
patch_size = 32,32,32
patch_overlap = 4,4,4
grid_sampler = tio.inference.GridSampler(
        subject,
        patch_size,
        patch_overlap
    )
patch_loader = torch.utils.data.DataLoader(
        grid_sampler, batch_size=4
    )
aggregator=tio.inference.GridAggregator(grid_sampler)
with torch.no_grad():
     for patches_batch in patch_loader:
            locations = patches_batch[tio.LOCATION]
            probabilities = model_dict.infer_batch(patches_batch)
            aggregator.add_batch(probabilities[0],locations)

When I print probabilities[0], where the predictions are saved, there are only NaN values. I have already checked and there are no NaN values in the test samples. I'm thinking that maybe I am loading the model incorrectly, since in validation during training it was working fine. I have also tested the loaded model on the validation set used during training and it does not work (even though it was working correctly during training). Any suggestions? Thanks in advance.