21 lines
399 B
Python
21 lines
399 B
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Fri Jan 19 16:10:39 2024
|
|
|
|
@author: ym
|
|
"""
|
|
import torch
|
|
from model.resnet_pre import resnet18
|
|
|
|
|
|
|
|
def main():
|
|
model_path = "best.pth"
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
model = resnet18().to(device)
|
|
model.load_state_dict(torch.load(model_path, map_location=device))
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |