diff --git a/C1_Browser-based-TF-JS/W4/assignment/C1_W4_Assignment.js b/C1_Browser-based-TF-JS/W4/assignment/C1_W4_Assignment.js
index e055f3cb..49c4d644 100755
--- a/C1_Browser-based-TF-JS/W4/assignment/C1_W4_Assignment.js
+++ b/C1_Browser-based-TF-JS/W4/assignment/C1_W4_Assignment.js
@@ -25,20 +25,24 @@ async function train() {
   // using ReLu activation functions where applicable.
   model = tf.sequential({
     layers: [
-        
-      // YOUR CODE HERE
-
+      tf.layers.flatten({inputShape: mobilenet.outputs[0].shape.slice(1)}),
+      tf.layers.dense({units: 100, activation: 'relu'}),
+      tf.layers.dense({units: 50, activation: 'relu'}),
+      tf.layers.dense({units: 5, activation: 'softmax'}) 
     ]
   });
     
    
   // Set the optimizer to be tf.train.adam() with a learning rate of 0.0001.
-  const optimizer = // YOUR CODE HERE
+  const optimizer = tf.train.adam(0.0001);
     
         
   // Compile the model using the categoricalCrossentropy loss, and
   // the optimizer you defined above.
-  model.compile(// YOUR CODE HERE);
+  model.compile({
+    optimizer: optimizer,
+    loss: 'categoricalCrossentropy'
+  });
  
   let loss = 0;
   model.fit(dataset.xs, dataset.ys, {
@@ -71,12 +75,10 @@ function handleButton(elem){
 			spockSamples++;
 			document.getElementById("spocksamples").innerText = "Spock samples:" + spockSamples;
 			break;
-            
-        // Add a case for lizard samples.
-        // HINT: Look at the previous cases.
-            
-        // YOUR CODE HERE
-		
+    case "4":
+      lizardSamples++;
+      document.getElementById("lizardsamples").innerText = "Lizard samples:" + lizardSamples;
+      break;
             
 	}
 	label = parseInt(elem.id);
@@ -108,12 +110,9 @@ async function predict() {
 		case 3:
 			predictionText = "I see Spock";
 			break;
-            
-        // Add a case for lizard samples.
-        // HINT: Look at the previous cases.
-            
-        // YOUR CODE HERE 
-	
+    case 4:
+      predictionText = "I see Lizard";
+      break;
             
 	}
 	document.getElementById("prediction").innerText = predictionText;
@@ -154,4 +153,4 @@ async function init(){
 }
 
 
-init();
\ No newline at end of file
+init();