I manage to solve the problem with following link How to load part of pre trained model? @apaszke post.