Skip to content

GitLab

Commit dcb5e4e4 authored by wangkx1's avatar wangkx1
Browse files

Update torch_verify.py

parent 32eb2157
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torch_verify.py
View file @ dcb5e4e4
...@@ -7,7 +7,7 @@ from PIL import Image ...@@ -7,7 +7,7 @@ from PIL import Image
# 加载预训练的ResNet50模型，并指定不使用GPU # 加载预训练的ResNet50模型，并指定不使用GPU
model = models.resnet50(pretrained=True) model = models.resnet50(pretrained=True)
model.eval() # 将模型设置为评估模式 model.eval() # 将模型设置为评估模式
device = torch.device("cpu") # 指定设备为CPU device = torch.device("cuda:0") # 指定设备为CPU
model.to(device) # 将模型转移到CPU上 model.to(device) # 将模型转移到CPU上
# 图片预处理 # 图片预处理
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment