diff --git a/Makefile b/Makefile
index f6c4b70..cecee96 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,5 @@
 CONFIG ?= release
-SRCS := source/main.cpp source/Rbm.cpp source/Layer.cpp source/Stack.cpp
+SRCS := source/main.cpp source/Rbm.cpp source/Layer.cpp source/Stack.cpp source/noise.c
 	
 	
 LIBS := -larmadillo -ljsoncpp
diff --git a/source/main.cpp b/source/main.cpp
index ea8710e..47e04a3 100644
--- a/source/main.cpp
+++ b/source/main.cpp
@@ -105,16 +105,19 @@ int main()
 	stack.save();
 
 	// Shake weights
-	stack.weightsInit(0.01);
+	stack.saveWeights();
+
 #else	
 	// Load project
 	stack.load();
+	
 	// Load weights
 	stack.loadWeights();
 #endif	
 	
 	// Train stack
-	stack.train(batch, 1000, 1000, &statusDisplay);
+	stack.weightsInit(0.01);
+	stack.train(batch, 100, 1000, &statusDisplay);
 
 	// Save weights
 	stack.saveWeights();
diff --git a/test.prj b/test.prj
index 3e15c0a..4a0f2f2 100644
--- a/test.prj
+++ b/test.prj
@@ -19,8 +19,7 @@
                   "momentum" : 0.5,
                   "numGibbs" : 1,
                   "weightDecay" : 0
-               },
-               "weightInit" : 0.01
+               }
             },
             "weights_file" : "Layer.0.weights.dat"
          }