Loading baseline with its weights throws the following error:
RuntimeError: Error(s) in loading state_dict for CNN:
size mismatch for dense.weight: copying a param with shape torch.Size([56, 64]) from checkpoint, the shape in current model is torch.Size([15, 64]).
size mismatch for dense.bias: copying a param with shape torch.Size([56]) from checkpoint, the shape in current model is torch.Size([15]).