How to save training weight checkpoint of model and continue training from last point in PyTorch?
I'm trying to save checkpoint weights of the trained model after a certain number of epochs and continue to train from that last checkpoint to another number of epochs using PyTorch
To achieve this I've written a script like below
To train the model:
def create_model():
  # load model from package
  model = smp.Unet(
      encoder_name=&quot;resnet152&quot;,        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
      encoder_weights='imagenet',           # use `imagenet` pre-trained weights for encoder initialization
      in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
      classes=2,                      # model output channels (number of classes in your dataset)
  )
  return model

model = create_model()
model.to(device)
learning_rate = 1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
epochs = 5

for epoch in range(epochs):
    print('Epoch: [{}/{}]'.format(epoch+1, epochs))

    # train set
    pbar = tqdm(train_loader)
    model.train()
    iou_logger = iouTracker()
    for batch in pbar:
        # load image and mask into device memory
        image = batch['image'].to(device)
        mask = batch['mask'].to(device)

        # pass images into model
        pred = model(image)
        # pred = checkpoint['model_state_dict']

        # get loss
        loss = criteria(pred, mask)

        # update the model
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # compute and display progress
        iou_logger.update(pred, mask)
        mIoU = iou_logger.get_mean()
        pbar.set_description('Loss: {0:1.4f} | mIoU {1:1.4f}'.format(loss.item(), mIoU))

    # development set
    pbar = tqdm(development_loader)
   
    model.eval()
    iou_logger = iouTracker()
    with torch.no_grad():
        for batch in pbar:
            # load image and mask into device memory
            image = batch['image'].to(device)
            mask = batch['mask'].to(device)

            # pass images into model
            pred = model(image)

            # get loss
            loss = criteria(pred, mask)
            
            # compute and display progress
            iou_logger.update(pred, mask)
            mIoU = iou_logger.get_mean()
            pbar.set_description('Loss: {0:1.4f} | mIoU {1:1.4f}'.format(loss.item(), mIoU))

# save model
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,}, '/content/drive/MyDrive/checkpoint.pt')

from this, I can save the model checkpoint file as checkpoint.pt for 5 epochs
To continue the training using the saved checkpoint weight file for another I wrote below script:
epochs = 5    
for epoch in range(epochs):
    print('Epoch: [{}/{}]'.format(epoch+1, epochs))

    # train set
    pbar = tqdm(train_loader)


    checkpoint = torch.load( '/content/drive/MyDrive/checkpoint.pt')
    print(checkpoint)
    

    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)

    
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint['epoch']
    loss = checkpoint['loss']
    model.train()
    iou_logger = iouTracker()
    for batch in pbar:
        # load image and mask into device memory
        image = batch['image'].to(device)
        mask = batch['mask'].to(device)

        # pass images into model
        pred = model(image)
        # pred = checkpoint['model_state_dict']

        # get loss
        loss = criteria(pred, mask)

        # update the model
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # compute and display progress
        iou_logger.update(pred, mask)
        mIoU = iou_logger.get_mean()
        pbar.set_description('Loss: {0:1.4f} | mIoU {1:1.4f}'.format(loss.item(), mIoU))

    # development set
    pbar = tqdm(development_loader)
   
    model.eval()
    iou_logger = iouTracker()
    with torch.no_grad():
        for batch in pbar:
            # load image and mask into device memory
            image = batch['image'].to(device)
            mask = batch['mask'].to(device)

            # pass images into model
            pred = model(image)

            # get loss
            loss = criteria(pred, mask)
            
            # compute and display progress
            iou_logger.update(pred, mask)
            mIoU = iou_logger.get_mean()
            pbar.set_description('Loss: {0:1.4f} | mIoU {1:1.4f}'.format(loss.item(), mIoU))

# save model
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,}, 'checkpoint.pt')

This throws error:
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-31-54f48c10531a&gt; in &lt;module&gt;()


---&gt; 14     model.load_state_dict(checkpoint['model_state_dict'])



/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1222         if len(error_msgs) &gt; 0:
   1223             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-&gt; 1224                                self.__class__.__name__, &quot;\n\t&quot;.join(error_msgs)))
   1225         return _IncompatibleKeys(missing_keys, unexpected_keys)
   1226 

RuntimeError: Error(s) in loading state_dict for DataParallel:
    Missing key(s) in state_dict: &quot;module.encoder.conv1.weight&quot;, &quot;module.encoder.bn1.weight&quot;, &quot;module.encoder.bn1.bias&quot;, &quot;module.encoder.bn1.running_mean&quot;, &quot;module.encoder.bn1.running_var&quot;, &quot;module.encoder.layer1.0.conv1.weight&quot;, &quot;module.encoder.layer1.0.bn1.weight&quot;, &quot;module.encoder.layer1.0.bn1.bias&quot;, &quot;module.encoder.layer1.0.bn1.running_mean&quot;, &quot;module.encoder.layer1.0.bn1.running_var&quot;, &quot;module.encoder.layer1.0.conv2.weight&quot;, &quot;module.encoder.layer1.0.bn2.weight&quot;, &quot;module.encoder.layer1.0.bn2.bias&quot;, &quot;module.encoder.layer1.0.bn2.running_mean&quot;, &quot;module.encoder.layer1.0.bn2.running_var&quot;, &quot;module.encoder.layer1.0.conv3.weight&quot;, &quot;module.encoder.layer1.0.bn3.weight&quot;, &quot;module.encoder.layer1.0.bn3.bias&quot;, &quot;module.encoder.layer1.0.bn3.running_mean&quot;, &quot;module.encoder.layer1.0.bn3.running_var&quot;, &quot;module.encoder.layer1.0.downsample.0.weight&quot;, &quot;module.encoder.layer1.0.downsample.1.weight&quot;, &quot;module.encoder.layer1.0.downsample.1.bias&quot;, &quot;module.encoder.layer1.0.downsample.1.running_mean&quot;, &quot;module.encoder.layer1.0.downsample.1.running_var&quot;, &quot;module.encoder.layer1.1.conv1.weight&quot;, &quot;module.encoder.layer1.1.bn1.weight&quot;, &quot;module.encoder.layer1.1.bn1.bias&quot;, &quot;module.encoder.layer1.1.bn1.running_mean&quot;, &quot;module.encoder.layer1.1.bn1.running_var&quot;, &quot;module.encoder.layer1.1.conv2.weight&quot;, &quot;module.encoder.layer1.1.bn2.weight&quot;, &quot;module.encoder.layer1.1.bn2.bias&quot;, &quot;module.encoder.layer1.1.bn2.running_mean&quot;, &quot;module.encoder.layer1.1.bn2.running_var&quot;, &quot;module.encoder.layer1.1.conv3.weight&quot;, &quot;module.encoder.layer...
    Unexpected key(s) in state_dict: &quot;encoder.conv1.weight&quot;, &quot;encoder.bn1.weight&quot;, &quot;encoder.bn1.bias&quot;, &quot;encoder.bn1.running_mean&quot;, &quot;encoder.bn1.running_var&quot;, &quot;encoder.bn1.num_batches_tracked&quot;, &quot;encoder.layer1.0.conv1.weight&quot;, &quot;encoder.layer1.0.bn1.weight&quot;, &quot;encoder.layer1.0.bn1.bias&quot;, &quot;encoder.layer1.0.bn1.running_mean&quot;, &quot;encoder.layer1.0.bn1.running_var&quot;, &quot;encoder.layer1.0.bn1.num_batches_tracked&quot;, &quot;encoder.layer1.0.conv2.weight&quot;, &quot;encoder.layer1.0.bn2.weight&quot;, &quot;encoder.layer1.0.bn2.bias&quot;, &quot;encoder.layer1.0.bn2.running_mean&quot;, &quot;encoder.layer1.0.bn2.running_var&quot;, &quot;encoder.layer1.0.bn2.num_batches_tracked&quot;, &quot;encoder.layer1.1.conv1.weight&quot;, &quot;encoder.layer1.1.bn1.weight&quot;, &quot;encoder.layer1.1.bn1.bias&quot;, &quot;encoder.layer1.1.bn1.running_mean&quot;, &quot;encoder.layer1.1.bn1.running_var&quot;, &quot;encoder.layer1.1.bn1.num_batches_tracked&quot;, &quot;encoder.layer1.1.conv2.weight&quot;, &quot;encoder.layer1.1.bn2.weight&quot;, &quot;encoder.layer1.1.bn2.bias&quot;, &quot;encoder.layer1.1.bn2.running_mean&quot;, &quot;encoder.layer1.1.bn2.running_var&quot;, &quot;encoder.layer1.1.bn2.num_batches_tracked&quot;, &quot;encoder.layer1.2.conv1.weight&quot;, &quot;encoder.layer1.2.bn1.weight&quot;, &quot;encoder.layer1.2.bn1.bias&quot;, &quot;encoder.layer1.2.bn1.running_mean&quot;, &quot;encoder.layer1.2.bn1.running_var&quot;, &quot;encoder.layer1.2.bn1.num_batches_tracked&quot;, &quot;encoder.layer1.2.conv2.weight&quot;, &quot;encoder.layer1.2.bn2.weight&quot;, &quot;encoder.layer1.2.bn2.bias&quot;, &quot;encoder.layer1.2.bn2.running_mean&quot;, &quot;encoder.layer1.2.bn2.running_var&quot;, &quot;encoder.layer1.2.bn2.num_batches_tracked&quot;, &quot;encoder.layer2.0.conv1.weight&quot;, &quot;encoder.layer...

What am I doing wrong? How can I fix this? Any help on this will be helpful.

this line:
model.load_state_dict(checkpoint['model_state_dict'])

should be like this:
model.load_state_dict(checkpoint)

