update code

zhangnan96 · Sep 4, 2020 · f081e72 · f081e72
1 parent 97aa5cd
commit f081e72
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -12,6 +12,9 @@ flower_data
 *.config
 *.gz
 *.onnx
+*.xml
+*.bin
+*.mapping
 checkpoint
 data
 VOCdevkit

diff --git a/pytorch_classification/Test5_resnet/predict.py b/pytorch_classification/Test5_resnet/predict.py
@@ -5,6 +5,8 @@
 import matplotlib.pyplot as plt
 import json
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 data_transform = transforms.Compose(
     [transforms.Resize(256),
      transforms.CenterCrop(224),
@@ -31,7 +33,7 @@
 model = resnet34(num_classes=5)
 # load model weights
 model_weight_path = "./resNet34.pth"
-model.load_state_dict(torch.load(model_weight_path))
+model.load_state_dict(torch.load(model_weight_path, map_location=device))
 model.eval()
 with torch.no_grad():
     # predict class