diff --git a/README.md b/README.md
index 8a9ed64..34a141a 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # SQLNet
 
-This repo provides an implementation of SQLNet and Seq2SQL neural networks for predicting SQL queries on [WikiSQL dataset](https://github.com/salesforce/WikiSQL). 
+This repo provides an implementation of SQLNet and Seq2SQL neural networks for predicting SQL queries on [WikiSQL dataset](https://github.com/salesforce/WikiSQL). The paper is available at [here](https://arxiv.org/abs/1711.04436).
 
 ## Citation
 
@@ -10,12 +10,10 @@ This repo provides an implementation of SQLNet and Seq2SQL neural networks for p
 
 ```
 @article{xu2017sqlnet,
-  author    = {Xiaojun Xu and
-               Chang Liu and
-               Dawn Song},
-  title     = {SQLNet: Generating Structured Queries from Natural Language Without
-               Reinforcement Learning},
-  year      = {2017},
+  title={SQLNet: Generating Structured Queries From Natural Language Without Reinforcement Learning},
+  author={Xu, Xiaojun and Liu, Chang and Song, Dawn},
+  journal={arXiv preprint arXiv:1711.04436},
+  year={2017}
 }
 ```
 
diff --git a/sqlnet/model/modules/seq2sql_condition_predict.py b/sqlnet/model/modules/seq2sql_condition_predict.py
index 64551ea..4bb34ba 100644
--- a/sqlnet/model/modules/seq2sql_condition_predict.py
+++ b/sqlnet/model/modules/seq2sql_condition_predict.py
@@ -83,7 +83,7 @@ def forward(self, x_emb_var, x_len, col_inp_var, col_name_len, col_len,
             else:
                 cur_inp = Variable(torch.from_numpy(init_inp))
             cur_h = decoder_hidden
-            while len(done_set) < B*4 and t < 100:
+            while len(done_set) < B and t < 100:
                 g_s, cur_h = self.cond_decoder(cur_inp, cur_h)
                 g_s_expand = g_s.unsqueeze(2)
 
diff --git a/train.py b/train.py
index f396068..ed0cab5 100644
--- a/train.py
+++ b/train.py
@@ -1,4 +1,3 @@
-### TODO: Need to be test yet!
 import json
 import torch
 from sqlnet.utils import *