diff --git a/insurance_qa_eval.py b/insurance_qa_eval.py
index e7e16c2..a822af0 100644
--- a/insurance_qa_eval.py
+++ b/insurance_qa_eval.py
@@ -165,11 +165,11 @@ class Evaluator:
                 question = self.padq([d['question']] * len(indices))
 
                 n_good = len(d['good'])
-                sims = model.predict([question, answers], batch_size=500).flatten()
+                sims = model.predict([question, answers])
                 r = rankdata(sims, method='max')
 
-                max_r = np.argmax(r)
-                max_n = np.argmax(r[:n_good])
+                max_r = np.argmax(sims)
+                max_n = np.argmax(sims[:n_good])
 
                 # print(' '.join(self.revert(d['question'])))
                 # print(' '.join(self.revert(self.answers[indices[max_r]])))
@@ -219,18 +219,18 @@ if __name__ == '__main__':
         'question_len': 50,
         'answer_len': 100,
         'n_words': 22353,  # len(vocabulary) + 1
-        'margin': 0.009,
+        'margin': 0.02,
 
         'training_params': {
             'save_every': 1,
             'batch_size': 20,
-            'nb_epoch': 10,
-            'validation_split': 0.2,
+            'nb_epoch': 50,
+            'validation_split': 0.1,
             'optimizer': Adam(clipnorm=1e-2),
         },
 
         'model_params': {
-            'n_embed_dims': 1000,
+            'n_embed_dims': 100,
             'n_hidden': 200,
 
             # convolution
@@ -240,8 +240,8 @@ if __name__ == '__main__':
             # recurrent
             'n_lstm_dims': 141,  # * 2
 
-            'initial_embed_weights': np.load('models/word2vec_1000_dim.h5'),
-            'similarity_dropout': 0.2,
+            'initial_embed_weights': np.load('models/word2vec_100_dim.h5'),
+            'similarity_dropout': 0.5,
         },
 
         'similarity_params': {
@@ -255,8 +255,8 @@ if __name__ == '__main__':
     evaluator = Evaluator(conf)
 
     ##### Define model ######
-    model = EmbeddingModel(conf)
-    optimizer = conf.get('training_params', dict()).get('optimizer', 'adam')
+    model = AttentionModel(conf)
+    optimizer = conf.get('training_params', dict()).get('optimizer', 'rmsprop')
     model.compile(optimizer=optimizer)
 
     # save embedding layer
@@ -271,7 +271,7 @@ if __name__ == '__main__':
 
     # evaluate mrr for a particular epoch
     evaluator.load_epoch(model, best_loss['epoch'])
-    # evaluator.load_epoch(model, 68)
+    # evaluator.load_epoch(model, 31)
     evaluator.get_mrr(model, evaluate_all=True)
     # for epoch in range(1, 100):
     #     print('Epoch %d' % epoch)
diff --git a/keras_models.py b/keras_models.py
index 3ab9686..5748834 100644
--- a/keras_models.py
+++ b/keras_models.py
@@ -3,13 +3,11 @@ from __future__ import print_function
 from abc import abstractmethod
 
 from keras.engine import Input
-from keras.layers import merge, Embedding, Dropout, Convolution1D, Lambda, Activation, LSTM, Dense, TimeDistributed, \
-    ActivityRegularization, constraints, regularizers
+from keras.layers import merge, Embedding, Dropout, Convolution1D, Lambda, LSTM, Dense, TimeDistributed, constraints
 from keras import backend as K
 from keras.models import Model
 
 import numpy as np
-from keras.regularizers import EigenvalueRegularizer
 
 from attention_lstm import AttentionLSTM
 
@@ -72,7 +70,7 @@ class LanguageModel:
 
         axis = lambda a: len(a._keras_shape) - 1
         dot = lambda a, b: K.batch_dot(a, b, axes=axis(a))
-        l2_norm = lambda a, b: K.sqrt(K.sum((a - b) ** 2, axis=axis(a), keepdims=True))
+        l2_norm = lambda a, b: K.sqrt(((a - b) ** 2).sum())
 
         if similarity == 'cosine':
             return lambda x: dot(x[0], x[1]) / K.sqrt(dot(x[0], x[0]) * dot(x[1], x[1]))
@@ -107,7 +105,7 @@ class LanguageModel:
             similarity = self.get_similarity()
             qa_model = merge([dropout(question_output), dropout(answer_output)],
                              mode=similarity, output_shape=lambda _: (None, 1))
-            self._qa_model = Model(input=[self.question, self.get_answer()], output=[qa_model])
+            self._qa_model = Model(input=[self.question, self.get_answer()], output=qa_model)
 
         return self._qa_model
 
@@ -132,8 +130,9 @@ class LanguageModel:
         y = np.zeros(shape=(x[0].shape[0],))
         return self.training_model.fit(x, y, **kwargs)
 
-    def predict(self, x, **kwargs):
-        return self.prediction_model.predict(x, **kwargs)
+    def predict(self, x):
+        assert self.prediction_model is not None and isinstance(self.prediction_model, Model)
+        return self.prediction_model.predict_on_batch(x)
 
     def save_weights(self, file_name, **kwargs):
         assert self.prediction_model is not None, 'Must compile the model before saving weights'
@@ -154,7 +153,7 @@ class EmbeddingModel(LanguageModel):
         weights = weights if weights is None else [weights]
         embedding = Embedding(input_dim=self.config['n_words'],
                               output_dim=self.model_params.get('n_embed_dims', 100),
-                              W_constraint=constraints.nonneg(),
+                              # W_constraint=constraints.nonneg(),
                               weights=weights,
                               mask_zero=True)
         question_embedding = embedding(question)
@@ -209,8 +208,8 @@ class ConvolutionModel(LanguageModel):
         answer_cnn = merge([cnn(answer_dense) for cnn in cnns], mode='concat')
 
         # maxpooling
-        maxpool = Lambda(lambda x: K.max(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
-        avepool = Lambda(lambda x: K.mean(x, axis=1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
+        maxpool = Lambda(lambda x: K.max(x, axis=-1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
+        avepool = Lambda(lambda x: K.mean(x, axis=-1, keepdims=False), output_shape=lambda x: (x[0], x[2]))
         question_pool = maxpool(question_cnn)
         answer_pool = maxpool(answer_cnn)