diff --git a/sm-model/main.py b/sm-model/main.py
index 2f6cefb..0e963ff 100644
--- a/sm-model/main.py
+++ b/sm-model/main.py
@@ -90,7 +90,7 @@ if __name__ == "__main__":
     
     # epoch related arguments
     ap.add_argument('--epochs', type=int, default=25)
-    ap.add_argument('--patience', type=int, default=3, help="if there is no appreciable change in model after <patience> epochs, then stop")
+    ap.add_argument('--patience', type=int, default=5, help="if there is no appreciable change in model after <patience> epochs, then stop")
     
     # debugging arguments
     ap.add_argument('--debugSingleBatch', action="store_true", help="will stop program after training 1 input batch")
@@ -116,17 +116,19 @@ if __name__ == "__main__":
     torch.set_num_threads(args.num_threads)
     
     trainer = Trainer(net, args.eta, args.mom, args.no_loss_reg)
+    trainer.load_input_data(args.dataset_folder, cache_file, 'train', 'clean-dev', 'clean-test')
 
     best_map = 0.0
     best_model = 0
     
     for i in range(args.epochs):
         logger.info('------------- Training epoch {} --------------'.format(i+1))        
-        train_accuracy = trainer.train(args.dataset_folder, 'train', args.batch_size, cache_file, args.debugSingleBatch)        
+        train_accuracy = trainer.train('train', args.batch_size, args.debugSingleBatch)        
         if args.debugSingleBatch: sys.exit(0)
-        dev_accuracy, dev_scores = trainer.test(args.dataset_folder, 'clean-dev', args.batch_size, cache_file)
+        dev_scores = trainer.test('clean-dev', args.batch_size)
+
         dev_map, dev_mrr = compute_map_mrr(args.dataset_folder, 'clean-dev', dev_scores)
-        logger.info("MAP {}\nMRR {}".format(dev_map, dev_mrr))
+        logger.info("------- MAP {}, MRR {}".format(dev_map, dev_mrr))
 
         if np.fabs(dev_map - best_map) > 1e-3:            
             best_model = i
@@ -143,11 +145,10 @@ if __name__ == "__main__":
 
     model = QAModel.load(args.dataset_folder, args.model_fname)
     evaluator = Trainer(model, args.eta, args.mom, args.no_loss_reg)
-    test_accuracy, test_scores = evaluator.test(args.dataset_folder, 'clean-test', args.batch_size, cache_file)
-
-    logger.info('Test set accuracy = {:.4f}'.format(test_accuracy))
-
+    evaluator.load_input_data(args.dataset_folder, cache_file, None, None, 'clean-test')
+    test_scores = evaluator.test('clean-test', args.batch_size)
+    
     map, mrr = compute_map_mrr(args.dataset_folder, 'clean-test', test_scores)
-    logger.info("MAP {}\nMRR {}".format(map, mrr))
+    logger.info("------- MAP {}, MRR {}".format(map, mrr))
 
     
diff --git a/sm-model/train.py b/sm-model/train.py
index 3ddab13..4c46ee4 100644
--- a/sm-model/train.py
+++ b/sm-model/train.py
@@ -35,7 +35,16 @@ class Trainer(object):
         #self.criterion = nn.NLLLoss()
         self.optimizer = optim.SGD(self.model.parameters(), lr=eta, momentum=mom, weight_decay=(0 if no_loss_reg else self.reg) ) 
 
+        self.datasets = {}
+        self.embeddings = {}
 
+    def load_input_data(self, dataset_root_folder, word_vectors_cache_file, train_set_folder, dev_set_folder, test_set_folder):
+        for set_folder in [train_set_folder, dev_set_folder, test_set_folder]:
+            if set_folder:
+                self.datasets[set_folder] = utils.read_in_dataset(dataset_root_folder, set_folder)
+                # NOTE: self.datasets[set_folder] = questions, sentences, labels, vocab, maxlen_q, maxlen_s, ext_feats 
+                self.embeddings[set_folder] = utils.load_cached_embeddings(word_vectors_cache_file, self.datasets[set_folder][3])
+        
     def regularize_loss(self, loss):       
         
         flattened_params = []
@@ -97,15 +106,12 @@ class Trainer(object):
         return torch.sum(y.data.long() == best)
 
 
-    def test(self, dataset_folder, set_folder, batch_size, word_vectors_cache_file):
+    def test(self, set_folder, batch_size):
         logger.info('----- Predictions on {} '.format(set_folder))
                 
-        questions, sentences, labels, vocab, maxlen_q, maxlen_s, ext_feats = \
-                utils.read_in_dataset(dataset_folder, set_folder)
-
-        # load word embeddings for training set vocab
-        word_vectors, vec_dim = utils.load_cached_embeddings(word_vectors_cache_file, vocab)
-
+        questions, sentences, labels, vocab, maxlen_q, maxlen_s, ext_feats = self.datasets[set_folder]
+        word_vectors, vec_dim = self.embeddings[set_folder]
+        
         self.model.eval()
 
         batch_size = 1
@@ -141,23 +147,19 @@ class Trainer(object):
             ypc += 1         
                
         # logger.info('{}_correct {}'.format(set_folder, total_correct))
-        logger.info('{}_loss {}'.format(set_folder, total_loss.data[0]))
+        # logger.info('{}_loss {}'.format(set_folder, total_loss.data[0]))
         logger.info('{} total {}'.format(set_folder, len(labels)))
         # logger.info('{}_loss = {:.4f}, acc = {:.4f}'.format( set_folder, total_loss.data[0]/len(labels), float(total_correct)/len(labels) ))
-        logger.info('{}_loss = {:.4f}'.format( set_folder, total_loss.data[0]/len(labels) ))
+        #logger.info('{}_loss = {:.4f}'.format( set_folder, total_loss.data[0]/len(labels) ))
 
-        return float(total_correct)/len(labels), y_pred
+        return y_pred
 
 
-    def train(self, dataset_folder, set_folder, batch_size, word_vectors_cache_file, debugSingleBatch):
+    def train(self, set_folder, batch_size, debugSingleBatch):
         train_start_time = time.time()
 
-        # read in training data
-        questions, sentences, labels, vocab, maxlen_q, maxlen_s, ext_feats = \
-                utils.read_in_dataset(dataset_folder, set_folder)
-
-        # load word embeddings for training set vocab
-        word_vectors, vec_dim = utils.load_cached_embeddings(word_vectors_cache_file, vocab)
+        questions, sentences, labels, vocab, maxlen_q, maxlen_s, ext_feats = self.datasets[set_folder]
+        word_vectors, vec_dim = self.embeddings[set_folder]
 
         # set model for training modep
         self.model.train()