diff --git a/index.js b/index.js
index 6ee0bb1..f97a345 100644
--- a/index.js
+++ b/index.js
@@ -1,6 +1,6 @@
 console.log("Hello Autoencoder 🚂");
 
-import * as tf from "@tensorflow/tfjs-node";
+import * as tf from "@tensorflow/tfjs-node-gpu";
 // import canvas from "canvas";
 // const { loadImage } = canvas;
 import Jimp from "jimp";
@@ -11,24 +11,25 @@ const W = 28;
 main();
 
 async function main() {
+
+  console.log(`Using backend: ${tf.getBackend()}`)
+
   // Build the model
-  const { decoderLayers, autoencoder } = buildModel();
+  const { decoder, autoencoder } = buildModel();
   // load all image data
-  const images = await loadImages(5100);
+  const images = await loadImages(7000);
   // train the model
-  const x_train = tf.tensor2d(images.slice(0, 5000));
-  await trainModel(autoencoder, x_train, 200);
+  const x_train = tf.tensor2d(images.slice(0, 5600));
+  await trainModel(autoencoder, x_train, 15);
   const saveResults = await autoencoder.save("file://public/model/");
 
-  console.log(autoencoder.summary());
-
   // const autoencoder = await tf.loadLayersModel("file://public/model/model.json");
   // test the model
   const x_test = tf.tensor2d(images.slice(5000));
   await generateTests(autoencoder, x_test);
 
   // Create a new model with just the decoder
-  const decoder = createDecoder(decoderLayers);
+  //const decoder = createDecoder(decoder);
   const saveDecoder = await decoder.save("file://public/decoder/model/");
 }
 
@@ -61,19 +62,10 @@ async function generateTests(autoencoder, x_test) {
   }
 }
 
-function createDecoder(decoderLayers) {
-  const decoder = tf.sequential();
-  for (let layer of decoderLayers) {
-    const newLayer = tf.layers.dense({
-      units: layer.units,
-      activation: layer.activation,
-      inputShape: [layer.kernel.shape[0]],
-    });
-    decoder.add(newLayer);
-    newLayer.setWeights(layer.getWeights());
-  }
-  // const learningRate = 0.000001;
-  // const optimizer = tf.train.adam(learningRate);
+function createDecoder(decoder) {
+
+   const learningRate = 0.0001;
+   const optimizer = tf.train.adam(learningRate, 0.0000001); // adam(learning_rate, decay)
   decoder.compile({
     optimizer: "adam",
     loss: "meanSquaredError",
@@ -82,87 +74,51 @@ function createDecoder(decoderLayers) {
 }
 
 function buildModel() {
-  const autoencoder = tf.sequential();
-  // Build the model
 
-  // Encoder
-  autoencoder.add(
-    tf.layers.dense({
-      units: 256,
-      inputShape: [W * W],
-      activation: "relu",
-    })
-  );
-  autoencoder.add(
-    tf.layers.dense({
-      units: 128,
-      activation: "relu",
-    })
-  );
-  autoencoder.add(
-    tf.layers.dense({
-      units: 64,
-      activation: "relu",
-    })
-  );
-  autoencoder.add(
-    tf.layers.dense({
-      units: 16,
-      activation: "relu",
-    })
-  );
-  autoencoder.add(
-    tf.layers.dense({
-      units: 4,
-      activation: "sigmoid",
-    })
-  );
-  // How do I start from here?
-  // Decoder
-
-  let decoderLayers = [];
-  decoderLayers.push(
-    tf.layers.dense({
-      units: 16,
-      activation: "relu",
-    })
-  );
-  decoderLayers.push(
-    tf.layers.dense({
-      units: 64,
-      activation: "relu",
-    })
-  );
-  decoderLayers.push(
-    tf.layers.dense({
-      units: 128,
-      activation: "relu",
-    })
-  );
-  decoderLayers.push(
-    tf.layers.dense({
-      units: 256,
-      activation: "relu",
-    })
-  );
-  decoderLayers.push(
-    tf.layers.dense({
-      units: W * W,
-      activation: "sigmoid",
-    })
-  );
-
-  for (let layer of decoderLayers) {
-    autoencoder.add(layer);
-  }
 
-  // const learningRate = 0.001;
-  // const optimizer = tf.train.adam(learningRate);
+// encoder
+  const encoder_input = tf.input({shape: [W*W], name: "encoder_input"});
+ // const l0 = tf.layers.flatten().apply(encoder_input);
+  const l1 = tf.layers.dense({units: 128, activation: "relu"}).apply(encoder_input);
+  const l2 = tf.layers.dense({units: 64, activation: "relu"}).apply(l1);
+  const l3 = tf.layers.dense({units: 16, activation: "relu"}).apply(l2);
+  const l4 = tf.layers.dense({units: 4, activation: "relu"}).apply(l3);
+  let encoded = tf.layers.dense({units: 2, activation: "relu", name: "encoder_output"}).apply(l4);
+  
+  let encoder = tf.model({inputs: encoder_input, outputs: encoded, name: "encoder"});
+  console.log(`Encoder Summary: ${encoder.summary()}`);
+
+  const decoder_input = tf.input({shape: [2]});
+  let decoder = tf.layers.dense({units: 4, activation: "relu", name: "decoder_input"}).apply(decoder_input);
+  const l6 = tf.layers.dense({units: 16, activation: "relu"}).apply(decoder);
+  const l7 = tf.layers.dense({units: 64, activation: "relu"}).apply(l6);
+  const l8 = tf.layers.dense({units: 128, activation: "relu"}).apply(l7);
+  let decoded = tf.layers.dense({units: 784, activation: "sigmoid", name: "decoder_output"}).apply(l8);
+
+  decoder = tf.model({inputs: decoder_input, outputs: decoded, name: "decoder"});
+  console.log(`Decoder Summary: ${decoder.summary()}`);
+
+
+  const auto = tf.input({shape: [W*W]});
+  encoded = encoder.apply(auto);
+  decoded = decoder.apply(encoded);
+
+  const autoencoder = tf.model({inputs: auto, outputs: decoded, name: "autoencoder"});
+  console.log(`Autoencoder Summary: ${autoencoder.summary()}`);
+
+ 
+  const learningRate = 0.0001;
+  const optimizer = tf.train.adam(learningRate, 0.000001); // adam(learning_rate, decay)
+  
   autoencoder.compile({
     optimizer: "adam",
     loss: "meanSquaredError",
   });
-  return { decoderLayers, autoencoder };
+  decoder.compile({
+    optimizer: "adam",
+    loss: "meanSquaredError"
+  });
+  return { decoder, autoencoder};
 }
 
 async function trainModel(autoencoder, x_train, epochs) {
@@ -171,6 +127,7 @@ async function trainModel(autoencoder, x_train, epochs) {
     batch_size: 32,
     shuffle: true,
     verbose: true,
+    validation_split: 0.1
   });
 }
 
@@ -179,7 +136,7 @@ async function loadImages(total) {
   for (let i = 0; i < total; i++) {
     const num = numeral(i).format("0000");
     const img = await Jimp.read(
-      `AutoEncoder_TrainingData/data/square${num}.png`
+      `data/square${num}.png`
     );
 
     let rawData = [];