How to visualize the feature map of each layer of the convolutional network

How do I replace the following code with my own network model?
########################################################
self.pretrained_model = models.vgg16(pretrained=True).features

1.Use register_forward_hook to get feature map



2. Use Captum or pytorch-cnn-visualizations to visualize the feature map