Fix prediction on CPU in DocReader.

duyphongdn1997 · Mar 12, 2019 · 1121aa3 · 1121aa3
1 parent d27180f
commit 1121aa3
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/drqa/reader/model.py b/drqa/reader/model.py
@@ -283,8 +283,7 @@ def predict(self, ex, candidates=None, top_n=1, async_pool=None):
             inputs = [e if e is None else e.cuda(non_blocking=True)
                       for e in ex[:5]]
         else:
-            inputs = [e if e is None else e.cuda(non_blocking=True)
-                      for e in ex[:5]]
+            inputs = [e for e in ex[:5]]
 
         # Run forward
         with torch.no_grad():