Commit 09f083dc authored by gd1551's avatar gd1551 Committed by GitHub

default tag count to 0 if not found

parent 5b0acfe0
...@@ -745,7 +745,7 @@ class EmbedderModel(nn.Module): ...@@ -745,7 +745,7 @@ class EmbedderModel(nn.Module):
D, I = D.squeeze(), I.squeeze() D, I = D.squeeze(), I.squeeze()
for i, id in enumerate(I): for i, id in enumerate(I):
tag = index[id] tag = index[id]
count = tag_count[tag] count = 0 if tag not in tag_count else tag_count[tag]
prob = D[i] prob = D[i]
results.append([tag, count, prob]) results.append([tag, count, prob])
......
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