pytorch 神经网络存取

import torch
from torch import nn

#trap
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x


tudui = Tudui()
torch.save(tudui, 'tudui_method.pth')
'''
after training
'''
model = torch.load('tudui_method.pth')
print(model)