#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/7 11:06
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : model_load.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
from torch import load
from torchvision.models import vgg16_bn
from model_save import MyNet

# load 1
model1 = load("vgg16_bn_net1.pth")
print(model1)
# load 2
vgg16_bn_net2 = vgg16_bn(pretrained=False)
vgg16_bn_net2.load_state_dict(load("vgg16_bn_net2.pth"))
print(vgg16_bn_net2)
# trap
my_net = load("my_net.pth")
print(my_net)