Issue with epoch count with repeated save/restore
🐛 Bug
I'm trying to save and restore the state of both a model and a pytorch-lightning trainer.
I suspect the epoch count is wrong because I'm not able to save and restore several times with the same max_epoch count.
Here's what I do:
Step 1: run model for max_epochs = 1. Save checkpoint (gets saved as epoch=0.ckpt)
Step 2: load previous checkpoint and rerun again with max_epochs = 1. No training is run (because 1 epoch was already run before). A checkpoint is saved again, however this is called epoch=1.ckpt.
Step 3: load checkpoint from step 2 and rerun again with max_epochs = 1. Training fails because it believes step 2 was run for 2 epochs (and here max is 1)
Output:
pytorch_lightning.utilities.exceptions.MisconfigurationException: 
            you restored a checkpoint with current_epoch=2
            but the Trainer(max_epochs=1)

Code below to reproduce.
What am I doing wrong? this should be a possible scenario right?
Thanks!
To Reproduce
Run code below 3 times from same location
import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
import pytorch_lightning as pl
from pathlib import Path

class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        # training_step defined the train loop. It is independent of forward
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train, val, _ = random_split(dataset, [5500, 500, len(dataset) - 6000])

def get_last_checkpoint(checkpoint_folder):
    if os.path.exists(checkpoint_folder):
        past_experiments = sorted(Path(checkpoint_folder).iterdir(), key=os.path.getmtime)

        for experiment in past_experiments[::-1]:
            experiment_folder = os.path.join(experiment, "checkpoints")
            if os.path.exists(experiment_folder):
                checkpoints = os.listdir(experiment_folder)

                if len(checkpoints):
                    checkpoints.sort()
                    path = os.path.join(experiment_folder, checkpoints[-1])
                    return path

    return None

chk = get_last_checkpoint('lightning_logs')
if chk is not None:
    print("loading from ", chk)
    autoencoder = LitAutoEncoder.load_from_checkpoint(chk)
else:
    autoencoder = LitAutoEncoder()

trainer = pl.Trainer(max_epochs=1, resume_from_checkpoint=chk)
trainer.fit(autoencoder, DataLoader(train), DataLoader(val))

Expected behavior
Should not increase epochs on second run. Should be able to load checkpoint and save unchanged (several times)
Environment

CUDA:
- GPU:
- Quadro P2000 with Max-Q Design
- available:         True
- version:           10.2
Packages:
- numpy:             1.19.2
- pyTorch_debug:     True
- pyTorch_version:   1.8.0.dev20201014
- pytorch-lightning: 1.0.2
- tqdm:              4.50.2
System:
- OS:                Windows
- architecture:
- 64bit
- WindowsPE
- processor:         Intel64 Family 6 Model 158 Stepping 10, GenuineIntel
- python:            3.7.9
- version:           10.0.17763

Additional context