diff --git a/models/multisql_predictor.py b/models/multisql_predictor.py
index a800419..b8ad65e 100644
--- a/models/multisql_predictor.py
+++ b/models/multisql_predictor.py
@@ -16,15 +16,15 @@ def __init__(self, N_word, N_h, N_depth, gpu, use_hs):
         self.gpu = gpu
         self.use_hs = use_hs
 
-        self.q_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
+        self.q_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2),
                 num_layers=N_depth, batch_first=True,
                 dropout=0.3, bidirectional=True)
 
-        self.hs_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
+        self.hs_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2),
                 num_layers=N_depth, batch_first=True,
                 dropout=0.3, bidirectional=True)
 
-        self.mkw_lstm = nn.LSTM(input_size=N_word, hidden_size=N_h/2,
+        self.mkw_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2),
                 num_layers=N_depth, batch_first=True,
                 dropout=0.3, bidirectional=True)
 
@@ -35,7 +35,7 @@ def __init__(self, N_word, N_h, N_depth, gpu, use_hs):
         self.multi_out_c = nn.Linear(N_h, N_h)
         self.multi_out = nn.Sequential(nn.Tanh(), nn.Linear(N_h, 1))
 
-        self.softmax = nn.Softmax() #dim=1
+        self.softmax = nn.Softmax(dim=1) #dim=1
         self.CE = nn.CrossEntropyLoss()
         self.log_softmax = nn.LogSoftmax()
         self.mlsml = nn.MultiLabelSoftMarginLoss()
@@ -91,7 +91,11 @@ def forward(self, q_emb_var, q_len, hs_emb_var, hs_len, mkw_emb_var, mkw_len):
 
     def loss(self, score, truth):
         data = torch.from_numpy(np.array(truth))
-        truth_var = Variable(data.cuda())
+        if self.gpu:
+            truth_var = Variable(data.cuda())
+        else:
+            truth_var = Variable(data)
+        truth_var = torch._cast_Long(truth_var)    
         loss = self.CE(score, truth_var)
 
         return loss