一、安装torchsummary库

pip install torchsummary

二、代码

import torchvision.models as models
from torchsummary import summary

model = models.AlexNet()
model.to('cuda')
summary(model,(3,224, 224))