diff --git a/chess_model_3_pieces.h5 b/chess_model_3_pieces.h5
new file mode 100644
index 00000000..76efcbcb
Binary files /dev/null and b/chess_model_3_pieces.h5 differ
diff --git a/tensor_classifier.py b/tensor_classifier.py
index 66faa558..f86fa945 100644
--- a/tensor_classifier.py
+++ b/tensor_classifier.py
@@ -1,4 +1,4 @@
-from __future__ import absolute_import, division, print_function, unicode_literals
+#from __future__ import absolute_import, division, print_function, unicode_literals
 
 import tensorflow as tf
 from tensorflow.python.keras import datasets, layers, models
@@ -10,7 +10,7 @@ import runner
 from main import find_occupied_squares
 from util import POSITION
 
-new_model = models.load_model('test_chess_model.h5')
+new_model = models.load_model('chess_model_3_pieces.h5')
 new_model.summary()
 
 #board = cv2.imread("whole_boards/boards_for_empty/board_1554286488.605142_rank_3.png")
@@ -44,7 +44,7 @@ square_3 = square_3 / 255.0
 square_4 = square_4 / 255.0
 
 
-pieces = ['knight', 'rook']
+pieces = ['knight', 'rook', 'bishop']
 
 for pos, square in occupied:
     square = cv2.cvtColor(square, cv2.COLOR_BGR2GRAY)
@@ -54,7 +54,7 @@ for pos, square in occupied:
     text_color = 255
     cv2.putText(square, f"{pos} {pieces[int(np.argmax(test))]}", (0, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
                 color=(text_color,) * 3, thickness=3)
-    cv2.imshow(f"{pos}", square)
+    cv2.imwrite("lel", square)
 
 cv2.waitKey(0)
 
diff --git a/tmp/tensor.py b/tmp/tensor.py
index 6240ba79..3ef49a3d 100644
--- a/tmp/tensor.py
+++ b/tmp/tensor.py
@@ -1,7 +1,6 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import tensorflow as tf
-from tensorflow.python.keras import datasets, layers, models
 import glob
 import numpy as np
 import cv2
@@ -58,18 +57,18 @@ test_img = np.array(test_img).reshape((len(test_img),width, height, 1))
 # Normalize pixel values to be between 0 and 1
 train_images, test_images = training_img / 255.0, test_img / 255.0
 
-model = models.Sequential()
-model.add(layers.Conv2D(256, (3, 3), activation='relu', input_shape=(width, height, 1)))
-model.add(layers.MaxPooling2D((2, 2)))
-model.add(layers.Conv2D(128, (3, 3), activation='relu'))
-model.add(layers.MaxPooling2D((2, 2)))
-model.add(layers.Conv2D(64, (3, 3), activation='relu'))
-model.add(layers.MaxPooling2D((2, 2)))
-model.add(layers.Conv2D(64, (3, 3), activation='relu'))
+model = tf.keras.models.Sequential()
+model.add(tf.keras.layers.Conv2D(256, (3, 3), activation='relu', input_shape=(width, height, 1)))
+model.add(tf.keras.layers.MaxPooling2D((2, 2)))
+model.add(tf.keras.layers.Conv2D(128, (3, 3), activation='relu'))
+model.add(tf.keras.layers.MaxPooling2D((2, 2)))
+model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
+model.add(tf.keras.layers.MaxPooling2D((2, 2)))
+model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
 
-model.add(layers.Flatten())
-model.add(layers.Dense(64, activation='relu'))
-model.add(layers.Dense(3, activation='softmax'))
+model.add(tf.keras.layers.Flatten())
+model.add(tf.keras.layers.Dense(64, activation='relu'))
+model.add(tf.keras.layers.Dense(3, activation='softmax'))
 
 model.summary()
 
@@ -77,7 +76,7 @@ model.compile(optimizer='adam',
               loss='sparse_categorical_crossentropy',
               metrics=['accuracy'])
 
-model.fit(train_images, training_labels, epochs=10)
+model.fit(train_images, training_labels, epochs=3)
 
 test_loss, test_acc = model.evaluate(test_images, test_labels_)