diff --git a/logic/models/run.go b/logic/models/run.go
index 5e8ae15..df88ce9 100644
--- a/logic/models/run.go
+++ b/logic/models/run.go
@@ -7,7 +7,6 @@ import (
 	"net/http"
 	"os"
 	"path"
-	"strconv"
 
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/utils"
 	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"
@@ -131,6 +130,8 @@ func handleRun(handle *Handle) {
         vi := 0
         var predictions = results[0].Value().([][]float32)[0]
 
+        fmt.Println(predictions)
+
         for i, v := range predictions {
             if v > vmax {
                 vi = i
@@ -140,9 +141,24 @@ func handleRun(handle *Handle) {
 
         os.RemoveAll(run_path)
 
+        rows, err := handle.Db.Query("select name from model_classes where model_id=$1 and class_order=$2;", model.Id, vi)
+        if err != nil { return Error500(err) }
+        if !rows.Next() {
+            LoadDefineTemplate(w, "/models/edit.html", "run-model-card", c.AddMap(AnyMap{
+                "Model": model,
+                "NotFound": true,
+                "Result": nil,
+            }))
+            return nil
+        }
+
+        var name string
+        if err = rows.Scan(&name); err != nil { return nil }
+
+
         LoadDefineTemplate(w, "/models/edit.html", "run-model-card", c.AddMap(AnyMap{
             "Model": model,
-            "Result": strconv.Itoa(vi),
+            "Result": name,
         }))
 		return nil
 	})
diff --git a/logic/models/train/train.go b/logic/models/train/train.go
index effc4c0..a2f6374 100644
--- a/logic/models/train/train.go
+++ b/logic/models/train/train.go
@@ -83,12 +83,6 @@ func generateCvs(handle *Handle, run_path string, model_id string)  (count int,
     if err != nil { return }
     defer data.Close()
 
-    type row struct {
-        path string
-        class_order int
-    }
-
-
     f, err := os.Create(path.Join(run_path, "train.csv"))
     if err != nil { return }
     defer f.Close()
diff --git a/views/models/edit.html b/views/models/edit.html
index 48a58a7..a1a1857 100644
--- a/views/models/edit.html
+++ b/views/models/edit.html
@@ -294,7 +294,11 @@
 
 {{ define "run-model-card" }}
     <form hx-headers='{"REQUEST-TYPE": "html"}' enctype="multipart/form-data" hx-post="/models/run" hx-swap="outerHTML">
-        {{ if .Result }}
+        {{ if .NotFound }}
+            <div>
+                The class was not found
+            </div>
+        {{ else if .Result }}
             <div>
                 Img Class: {{.Result}}
             </div>
diff --git a/views/py/python_model_template.py b/views/py/python_model_template.py
index f33b23f..8fd087e 100644
--- a/views/py/python_model_template.py
+++ b/views/py/python_model_template.py
@@ -34,7 +34,7 @@ def pathToLabel(path):
 def decode_image(img):
   # channels were reduced to 1 since image is grayscale
   # TODO chnage channel number based if grayscale
-  img = tf.io.decode_png(img, channels=1)
+  img = tf.io.decode_png(img, channels=3)
 
   return tf.image.resize(img, image_size)