Fix bug writing predictions in run_squad_pytorch

38f740a1 · VictorSanh · ee29871f · 38f740a1
Commit 38f740a1 authored Nov 02, 2018 by VictorSanh
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

run_squad_pytorch.py run_squad_pytorch.py +2 -2

No files found.
--- a/run_squad_pytorch.py
+++ b/run_squad_pytorch.py
@@ -910,8 +910,8 @@ def main():
            #end_logits = [x.item() for x in end_logits]
            end_logits = [x.view(-1).detach().cpu().numpy() for x in end_logits]
            for idx, i in enumerate(unique_id):
-                s = start_logits[idx]
-                e = end_logits[idx]
+                s = [float(x) for x in start_logits[idx]]
+                e = [float(x) for x in end_logits[idx]]
                all_results.append(
                    RawResult(
                        unique_id=i,