Classifying Flowers With Transfer Learning

Saving the Model (Ckeckpoint)

#Extract the class_to_idx transformation 
model.class_to_idx = train_data.class_to_idx

#Put the model in CPU mode to allow predictions without having CUDA
model.to('cpu')

#Create the checkpoint
checkpoint = {'input_size':25088,
              'output_size':102,
              'hidden_layers':[each.out_features for each in classifier.hidden_layers],
              'drop_p':0.2,
              'state_dict':model.state_dict(),
              'class_to_idx': model.class_to_idx}

filepath = 'checkpoint.pth'

#Save the model
torch.save(checkpoint, filepath)