diff --git a/tensor_classifier.py b/tensor_classifier.py
index 8a28ea10..66faa558 100644
--- a/tensor_classifier.py
+++ b/tensor_classifier.py
@@ -1,7 +1,6 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import tensorflow as tf
-import tensorflow_datasets as tfds
 from tensorflow.python.keras import datasets, layers, models
 import glob
 import numpy as np
diff --git a/tmp/tensor.py b/tmp/tensor.py
index 432fbe38..6240ba79 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
-import tensorflow_datasets as tfds
 from tensorflow.python.keras import datasets, layers, models
 import glob
 import numpy as np
@@ -30,11 +29,21 @@ for _ in range(10):
         training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
         training_labels.append(0)
 
+for _ in range(10):
+    for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[:-50]:
+        training_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
+        training_labels.append(2)
+
 for _ in range(5):
     for filename in glob.glob(f"../training_images/rook/*_square/*.png")[-50:]:
         test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
         test_labels_.append(1)
 
+for _ in range(5):
+    for filename in glob.glob(f"../training_images/bishop/*_square/*.png")[-50:]:
+        test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
+        test_labels_.append(2)
+
 for _ in range(5):
     for filename in glob.glob(f"../training_images/knight/*_square/*.png")[-50:]:
         test_img.append(cv2.cvtColor(cv2.imread(filename), cv2.COLOR_BGR2GRAY))
@@ -50,7 +59,9 @@ test_img = np.array(test_img).reshape((len(test_img),width, height, 1))
 train_images, test_images = training_img / 255.0, test_img / 255.0
 
 model = models.Sequential()
-model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(width, height, 1)))
+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)))
@@ -58,7 +69,7 @@ model.add(layers.Conv2D(64, (3, 3), activation='relu'))
 
 model.add(layers.Flatten())
 model.add(layers.Dense(64, activation='relu'))
-model.add(layers.Dense(2, activation='softmax'))
+model.add(layers.Dense(3, activation='softmax'))
 
 model.summary()
 
@@ -66,12 +77,12 @@ model.compile(optimizer='adam',
               loss='sparse_categorical_crossentropy',
               metrics=['accuracy'])
 
-model.fit(train_images, training_labels, epochs=5)
+model.fit(train_images, training_labels, epochs=10)
 
 test_loss, test_acc = model.evaluate(test_images, test_labels_)
 
 print(test_acc)
 
 # Save entire model to a HDF5 file
-model.save('test_chess_model.h5')
+model.save('chess_model_3_pieces.h5')