diff --git a/loader.py b/loader.py
index c31c7e1..33e204d 100644
--- a/loader.py
+++ b/loader.py
@@ -108,7 +108,7 @@ def train(samples):
     # NOTE: Automatically continue training the existing model
     # Otherwise, starts a new model if there is no model file found.
     cnn = train_model(
-      path_model if os.path.isfile(path_model + '0.npz') else None,
+      path_model if os.path.isfile(path_model) else None,
       trainsetX,
       trainsetY,
       validsetX,
diff --git a/model/cnn.py b/model/cnn.py
index de23041..51c144d 100644
--- a/model/cnn.py
+++ b/model/cnn.py
@@ -12,7 +12,7 @@
 import lasagne
 import pickle
 from lasagne import layers
-from lasagne.updates import adagrad, adadelta
+from lasagne.updates import adagrad, adadelta, apply_momentum
 from lasagne.objectives import *
 from nolearn.lasagne import NeuralNet
 from nolearn.lasagne import visualize
@@ -80,6 +80,7 @@ def train(self,X,y,X_,y_,batch_size=1000,num_epochs=100,learn_rate=[0.001,0.001,
     loss   = [T.sqrt(T.mean((output[i] - outputy[i])**2)) for i in range(len(self.nets))]
     params = [layers.get_all_params(n) for n in self.nets]
     update = [adadelta(loss[i], params[i]) for i in range(len(self.nets))]
+    update = [apply_momentum(update[i], params[i], momentum=0.8) for i in range(len(self.nets))]
 
     print(colored('...Preparing training functions','green'))
     models  = [theano.function(
diff --git a/model/core.py b/model/core.py
index 0804abe..f5a4538 100644
--- a/model/core.py
+++ b/model/core.py
@@ -25,8 +25,8 @@
 def train_model(path, X, y, X_, y_, image_dim, final_vec_dim, epoch, batch_size):
 
   # Create a new CNN
-  if path is None:
-    print(colored('Creating a new CNN.','green'))
+  if path is None or not os.path.isfile(path + '0'):
+    print(colored('Creating a new CNN.','cyan'))
     cnn = CNN(image_dim, final_vec_dim) 
   else:
     print(colored('Loading model from : {0}'.format(path),'green'))
diff --git a/setup-ec2-and-run b/setup-ec2-and-run
index 583a5f5..457a482 100644
--- a/setup-ec2-and-run
+++ b/setup-ec2-and-run
@@ -17,7 +17,7 @@ mkdir -p data/out
 mkdir -p data/unfiltered
 sudo pip3 install -r requirements.txt 
 cd data/
-wget http://starcolon.com/data/train.tar.gz
+wget http://starcolon.com/data/dataset.tar.gz
 tar -xvf train.tar.gz 
 rm raw/._*
 rm train.tar.gz
diff --git a/vision/core.py b/vision/core.py
index a0eb24c..f7cf651 100644
--- a/vision/core.py
+++ b/vision/core.py
@@ -9,7 +9,7 @@
 from .filters import *
 
 def get_sample_dim(): # n
-  return 128 # total vector dimensionality : n^2 * 3
+  return 640 # total vector dimensionality : n^2 * 3
 
 def generate_filtered(sample,max_combination=5):
   # Make transformation functions