Loading a saved model

[cs-mshah]

To resume training I tried to load a saved model from checkpoint, but it ran into the issue where the expected state_dict had module as a prefix to the keys as compared to the saved state_dict.

This is similar to the following issue: pytorch forum, github issue.

I modified the loading code in the following way by adding strict=False for now:

start_epoch = 0
if args.resume:
    assert os.path.isfile(
        args.resume), "Error: no checkpoint directory found!"
    args.out = os.path.dirname(args.resume)
    checkpoint = torch.load(args.resume)
    best_acc = checkpoint['best_acc']
    start_epoch = checkpoint['epoch']
    model.load_state_dict(checkpoint['state_dict'], strict=False)
    optimizer.load_state_dict(checkpoint['optimizer'])
    ema_model.load_state_dict(checkpoint['ema_state_dict'], strict=False)