pip install torchsummary

# -*- coding: utf-8 -*-
"""
Created on Thu Jul  5 14:22:03 2018

@author: guokai_liu
"""

import torch
from torchvision import models
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)

summary(vgg, (3, 224, 224))A