Unexpected key(s) in state_dict Error when calling Trainer.test
Dear all,
I have a trainer
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from pytorch_lightning import LightningModule
from torch.nn import functional as F
from pytorch_lightning.metrics.functional import accuracy, f1_score, auroc


class TraversabilityModule(LightningModule):

    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

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

    def get_metrics(self, pred, y):
        return {
            'accuracy': accuracy(pred, y, num_classes=2),
            'f1': f1_score(pred, y, num_classes=2),
            # 'roc_auc': auroc(pred, y)
        }

    def step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        pred = torch.argmax(y_hat, dim=1)
        metrics = self.get_metrics(pred, y)
        loss = F.cross_entropy(y_hat, y)
        return loss, metrics

    def training_step(self, batch, batch_idx):
        loss, metrics = self.step(batch, batch_idx)
        return {'loss': loss, 'log': metrics }

    def validation_step(self, batch, batch_idx):
        loss, metrics = self.step(batch, batch_idx)
        return {'val_loss': loss, 'log': metrics }

    def validation_epoch_end(self, outputs):
        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
        acc_mean = torch.stack([x['log']['accuracy'] for x in outputs]).mean()
        f1_mean = torch.stack([x['log']['f1'] for x in outputs]).mean()
        # roc_auc_mean = torch.stack([x['log']['roc_auc'] for x in outputs]).mean()

        return {
            'val_loss': val_loss_mean,
            'val_f1': f1_mean,
            'progress_bar': {'f1': f1_mean},
            'log': {
                'val_accuracy': acc_mean,
                'val_f1': f1_mean,
                # 'roc_auc': roc_auc_mean
                
                }
            }

    def test_step(self, batch, batch_idx):
        loss, metrics = self.step(batch, batch_idx)
        return {'test_loss': loss, 'log': metrics}

    def test_epoch_end(self, outputs):
        val_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean()
        acc_mean = torch.stack([x['log']['accuracy'] for x in outputs]).mean()
        f1_mean = torch.stack([x['log']['f1'] for x in outputs]).mean()
        # roc_auc_mean = torch.stack([x['log']['roc_auc'] for x in outputs]).mean()

        return {
            'test_loss': val_loss_mean,
            'test_f1': f1_mean,
            'progress_bar': {'f1': f1_mean},
            'log': {
                'test_accuracy': acc_mean,
                'test_f1': f1_mean,
                # 'roc_auc': roc_auc_mean
                
                }
            }

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        scheduler = {'scheduler': ReduceLROnPlateau(optimizer, verbose=True),
                     'monitor': 'val_f1'}

        return [optimizer], [scheduler]
And I have a training loop
module = TraversabilityModule(model)

    trainer = pl.Trainer(
        gpus=1,
        max_epochs=params['epoches'],
        logger=comet_logger,
        checkpoint_callback=ModelCheckpoint(
            monitor='val_f1',
            filepath=project.checkpoint_dir / params['model'] / 'model',
            ),
        early_stop_callback=EarlyStopping(monitor='val_f1',
                                          patience=15))
    trainer.fit(module, train_dl, val_dl)
    
    trainer.test(test_dataloaders=test_dl)
I get the following error when I try to validate the test set (at trainer.test)
RuntimeError: Error(s) in loading state_dict for TraversabilityModule:
        Unexpected key(s) in state_dict: "model.encoder.gate.0.weight", "model.encoder.gate.1.weight", "model.encoder.gate.1.bias", "model.encoder.gate.1.running_mean", "model.encoder.gate.1.running_var", "model.encoder.gate.1.num_batches_tracked", "model.encoder.layers.0.layer.0.shortcut.0.weight", "model.encoder.layers.0.layer.0.shortcut.1.weight", "model.encoder.layers.0.layer.0.shortcut.1.bias", "model.encoder.layers.0.layer.0.shortcut.1.running_mean", "model.encoder.layers.0.layer.0.shortcut.1.running_var", "model.encoder.layers.0.layer.0.shortcut.1.num_batches_tracked", "model.encoder.layers.0.layer.0.convs.0.weight", "model.encoder.layers.0.layer.0.convs.0.bias", "model.encoder.layers.0.layer.0.convs.0.running_mean", "model.encoder.layers.0.layer.0.convs.0.running_var", "model.encoder.layers.0.layer.0.convs.0.num_batches_tracked", "model.encoder.layers.0.layer.0.convs.2.weight", "model.encoder.layers.0.layer.0.convs.3.weight", "model.encoder.layers.0.layer.0.convs.3.bias", "model.encoder.layers.0.layer.0.convs.3.running_mean", "model.encoder.layers.0.layer.0.convs.3.running_var", "model.encoder.layers.0.layer.0.convs.3.num_batches_tracked", "model.encoder.layers.0.layer.0.convs.5.weight", "model.encoder.layers.0.layer.0.se.se.0.weight", "model.encoder.layers.0.layer.0.se.se.0.bias", "model.encoder.layers.0.layer.0.se.se.2.weight", "model.encoder.layers.0.layer.0.se.se.2.bias", "model.encoder.layers.1.layer.0.shortcut.0.weight", "model.encoder.layers.1.layer.0.shortcut.1.weight", "model.encoder.layers.1.layer.0.shortcut.1.bias", "model.encoder.layers.1.layer.0.shortcut.1.running_mean", "model.encoder.layers.1.layer.0.shortcut.1.running_var", "model.encoder.layers.1.layer.0.shortcut.1.num_batches_tracked", "model.encoder.layers.1.layer.0.convs.0.weight", "model.encoder.layers.1.layer.0.convs.0.bias", "model.encoder.layers.1.layer.0.convs.0.running_mean", "model.encoder.layers.1.layer.0.convs.0.running_var", "model.encoder.layers.1.layer.0.convs.0.num_batches_tracked", "model.encoder.layers.1.layer.0.convs.2.weight", "model.encoder.layers.1.layer.0.convs.3.weight", "model.encoder.layers.1.layer.0.convs.3.bias", "model.encoder.layers.1.layer.0.convs.3.running_mean", "model.encoder.layers.1.layer.0.convs.3.running_var", "model.encoder.layers.1.layer.0.convs.3.num_batches_tracked", "model.encoder.layers.1.layer.0.convs.5.weight", "model.encoder.layers.1.layer.0.se.se.0.weight", "model.encoder.layers.1.layer.0.se.se.0.bias", "model.encoder.layers.1.layer.0.se.se.2.weight", "model.encoder.layers.1.layer.0.se.se.2.bias", "model.encoder.layers.2.layer.0.shortcut.0.weight", "model.encoder.layers.2.layer.0.shortcut.1.weight", "model.encoder.layers.2.layer.0.shortcut.1.bias", "model.encoder.layers.2.layer.0.shortcut.1.running_mean", "model.encoder.layers.2.layer.0.shortcut.1.running_var", "model.encoder.layers.2.layer.0.shortcut.1.num_batches_tracked", "model.encoder.layers.2.layer.0.convs.0.weight", "model.encoder.layers.2.layer.0.convs.0.bias", "model.encoder.layers.2.layer.0.convs.0.running_mean", "model.encoder.layers.2.layer.0.convs.0.running_var", "model.encoder.layers.2.layer.0.convs.0.num_batches_tracked", "model.encoder.layers.2.layer.0.convs.2.weight", "model.encoder.layers.2.layer.0.convs.3.weight", "model.encoder.layers.2.layer.0.convs.3.bias", "model.encoder.layers.2.layer.0.convs.3.running_mean", "model.encoder.layers.2.layer.0.convs.3.running_var", "model.encoder.layers.2.layer.0.convs.3.num_batches_tracked", "model.encoder.layers.2.layer.0.convs.5.weight", "model.encoder.layers.2.layer.0.se.se.0.weight", "model.encoder.layers.2.layer.0.se.se.0.bias", "model.encoder.layers.2.layer.0.se.se.2.weight", "model.encoder.layers.2.layer.0.se.se.2.bias", "model.decoder.decoder.weight", "model.decoder.decoder.bias". 

They are clearly the weights in my model (a variant of ResNet). The model path exists and the model is stored there.
Thank you.
Best Regards,
Francesco