Make Kim CNN ONNX-exportable (#136)

* Kim CNN - only set embedding for corresponding mode

* Kim CNN ONNX Export

* Specify dummy ONNX input size from command line
This commit is contained in:
Michael Tu
2018-08-04 17:30:24 -04:00
committed by GitHub
parent 82bf90f4bb
commit 8a00f9cdcd
6 changed files with 37 additions and 19 deletions
+1 -1
View File
@@ -13,7 +13,7 @@ class SSTEvaluator(Evaluator):
total_loss = 0
for batch_idx, batch in enumerate(self.data_loader):
scores = self.model(batch)
scores = self.model(batch.text)
n_dev_correct += (
torch.max(scores, 1)[1].view(batch.label.size()).data == batch.label.data).sum().item()
total_loss += F.cross_entropy(scores, batch.label, size_average=False).item()