Merge pull request Shawn1993#11 from SeanRosario/master

Fixed cuda tensor bug during prediction

Author: Shawn1993

main.py

@@ -100,15 +100,18 @@ def mr(text_field, label_field, **kargs):
 
 # train or predict
 if args.predict is not None:
-    label = train.predict(args.predict, cnn, text_field, label_field)
-    print('\n[Text]  {}[Label] {}\n'.format(args.predict, label))
+    label = train.predict(args.predict, cnn, text_field, label_field, args.cuda)
+    print('\n[Text]  {}\n[Label] {}\n'.format(args.predict, label))
 elif args.test :
     try:
         train.eval(test_iter, cnn, args) 
     except Exception as e:
         print("\nSorry. The test dataset doesn't  exist.\n")
 else :
     print()
-    train.train(train_iter, dev_iter, cnn, args)
-    
+    try:
+        train.train(train_iter, dev_iter, cnn, args)
+    except KeyboardInterrupt:
+        print('-' * 89)
+        print('Exiting from training early')
 

train.py

@@ -74,14 +74,16 @@ def eval(data_iter, model, args):
                                                                        size))
 
 
-def predict(text, model, text_field, label_feild):
+def predict(text, model, text_field, label_feild, cuda_flag):
     assert isinstance(text, str)
     model.eval()
     # text = text_field.tokenize(text)
     text = text_field.preprocess(text)
     text = [[text_field.vocab.stoi[x] for x in text]]
     x = text_field.tensor_type(text)
     x = autograd.Variable(x, volatile=True)
+    if cuda_flag:
+        x =x.cuda()
     print(x)
     output = model(x)
     _, predicted = torch.max(output, 1)