mirror of
https://github.com/wassname/Castor.git
synced 2026-09-11 11:51:23 +08:00
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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user