lichen/ieemoo-ai-isempty
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update ieemoo-ai-isempty.py.

Browse Source
This commit is contained in:
Brainway
2022-10-18 03:59:42 +00:00
committed by Gitee
parent 952d7e3a3c
commit 14558af47b
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
ieemoo-ai-isempty.py
View File

@ -82,7 +82,7 @@ class Predictor(object):
# self.model = torch.load(self.args.pretrained_model)
# else:
# self.model = torch.load(self.args.pretrained_model,map_location='cpu')
self.model = torch.load(self.args.pretrained_model)
self.model = torch.load(self.args.pretrained_model,map_location=torch.device('cpu'))
self.model.eval()
self.model.to("cuda")
@ -95,8 +95,8 @@ class Predictor(object):
else:
with torch.no_grad():
x = self.test_transform(img_data)
if torch.cuda.is_available():
x = x.cuda()
# if torch.cuda.is_available():
# x = x.cuda()
part_logits = self.model(x.unsqueeze(0))
probs = torch.nn.Softmax(dim=-1)(part_logits)
topN = torch.argsort(probs, dim=-1, descending=True).tolist()