How to modify the final FC layer based on the torch.model

You can get the idea from this post https://discuss.pytorch.org/t/how-to-perform-finetuning-in-pytorch/419