diff --git a/m.py b/m.py
index 3628340..f012fe7 100644
--- a/m.py
+++ b/m.py
@@ -35,26 +35,26 @@ def generate_start(i):
     return kl.MaxPooling2D((3, 3))(r5)
 
 def generate_type_a(i):
-    a11 = conv(i, CONV_SIZE, (1, 1), (3, 3))
+    a11 = conv(i, CONV_SIZE, (1, 1))
 
     a21 = avg(i, (3, 3))
     a22 = conv(a21, CONV_SIZE, (1, 1))
 
     a31 = conv(i, CONV_SIZE, (1, 1))
-    a32 = conv(a31, CONV_SIZE, (3, 3), (3, 3))
+    a32 = conv(a31, CONV_SIZE, (3, 3))
 
     a41 = conv(i, CONV_SIZE, (1, 1))
     a42 = conv(a41, CONV_SIZE, (3, 3))
-    a43 = conv(a42, CONV_SIZE, (3, 3), (3, 3))
+    a43 = conv(a42, CONV_SIZE, (3, 3))
 
     return kl.Concatenate()([a11, a22, a32, a43])
 
 def generate_ab_bridge(i):
     ab11 = conv(i, CONV_SIZE, (1, 1))
     ab12 = conv(ab11, CONV_SIZE, (3, 3))
-    ab13 = conv(ab12, CONV_SIZE, (3, 3), (3, 3))
+    ab13 = conv(ab12, CONV_SIZE, (3, 3))
 
-    ab21 = conv(i, CONV_SIZE, (3, 3), (3, 3))
+    ab21 = conv(i, CONV_SIZE, (3, 3))
 
     ab31 = max(i, (3, 3))
 
@@ -65,13 +65,13 @@ def generate_type_b(i):
     b12 = conv(b11, CONV_SIZE, (7, 1))
     b13 = conv(b12, CONV_SIZE, (1, 7))
     b14 = conv(b13, CONV_SIZE, (7, 1))
-    b15 = conv(b14, CONV_SIZE, (1, 7), (3, 3))
+    b15 = conv(b14, CONV_SIZE, (1, 7))
 
     b21 = conv(i, CONV_SIZE, (1, 1))
     b22 = conv(b21, CONV_SIZE, (1, 7))
-    b23 = conv(b22, CONV_SIZE, (7, 1), (3, 3))
+    b23 = conv(b22, CONV_SIZE, (7, 1))
 
-    b31 = conv(i, CONV_SIZE, (1, 1), (3, 3))
+    b31 = conv(i, CONV_SIZE, (1, 1))
 
     b41 = avg(i, (3, 3))
     b42 = conv(b41, CONV_SIZE, (1, 1))
@@ -82,10 +82,10 @@ def generate_bc_bridge(i):
     bc11 = conv(i, CONV_SIZE, (1, 1))
     bc12 = conv(bc11, CONV_SIZE, (7, 1))
     bc13 = conv(bc12, CONV_SIZE, (1, 7))
-    bc14 = conv(bc13, CONV_SIZE, (3, 3), (3, 3))
+    bc14 = conv(bc13, CONV_SIZE, (3, 3))
 
     bc21 = conv(i, CONV_SIZE, (1, 1))
-    bc22 = conv(bc21, CONV_SIZE, (1, 1), (3, 3))
+    bc22 = conv(bc21, CONV_SIZE, (1, 1))
 
     bc31 = max(i, (3, 3))
 
@@ -97,23 +97,24 @@ def generate_aux(i):
     u3 = kl.Flatten()(u2)
     u4 = kl.Dropout(DRPT_RATE)(u3)
     u5 = kl.Dense(MPPT_SIZE)(u4)
+    u6 = kl.Dense(TOUT_AMNT)(u5)
 
-    return kl.Dense(TOUT_AMNT)(u5)
+    return kl.Activation('softmax')(u6)
 
 def generate_type_c(i):
     c11 = conv(i, CONV_SIZE, (1, 1))
     c12 = conv(c11, CONV_SIZE, (3, 3))
-    c13 = conv(c12, CONV_SIZE, (1, 3), (3, 3))
-    c14 = conv(c12, CONV_SIZE, (3, 1), (3, 3))
+    c13 = conv(c12, CONV_SIZE, (1, 3))
+    c14 = conv(c12, CONV_SIZE, (3, 1))
 
     c21 = conv(i, CONV_SIZE, (1, 1))
-    c22 = conv(c21, CONV_SIZE, (1, 3), (3, 3))
-    c23 = conv(c21, CONV_SIZE, (3, 1), (3, 3))
+    c22 = conv(c21, CONV_SIZE, (1, 3))
+    c23 = conv(c21, CONV_SIZE, (3, 1))
 
     c31 = max(i, (3, 3))
     c32 = conv(c31, CONV_SIZE, (1, 1))
 
-    c41 = conv(i, CONV_SIZE, (1, 1), (3, 3))
+    c41 = conv(i, CONV_SIZE, (1, 1))
 
     return kl.Concatenate()([c14, c23, c32, c41])
 
@@ -123,11 +124,12 @@ def generate_finish(i):
     f3 = kl.Flatten()(f2)
     f4 = kl.Dropout(DRPT_RATE)(f3)
     f5 = kl.Dense(MPPT_SIZE)(f4)
+    f6 = kl.Dense(TOUT_AMNT)(f5)
 
-    return kl.Dense(TOUT_AMNT)(f5)
+    return kl.Activation('softmax')(f6)
 
 
-gi = kl.Input((2048, 2048, 3))
+gi = kl.Input((300, 300, 3))
 
 ds = generate_start(gi)
 
@@ -148,7 +150,9 @@ for _ in range(1):
 
 go = generate_finish(go)
 
-mod = k.Model(inputs = gi, outputs = [go, uo])
+mod = k.Model(inputs = gi, outputs = go)
 
 mod.summary()
-mod.compile(optimizer = ko.Lion(learning_rate = 0.001), metrics = ['accuracy', 'val_accuracy'])
+mod.compile(optimizer = ko.Lion(learning_rate = 0.001),
+            loss = 'categorical_crossentropy',
+            metrics = ['accuracy'])
diff --git a/train_l.py b/train_l.py
new file mode 100644
index 0000000..02cdbe9
--- /dev/null
+++ b/train_l.py
@@ -0,0 +1,15 @@
+import tensorflow as tf
+
+ds_train, ds_valid = tf.keras.preprocessing.image_dataset_from_directory(
+    '../dataset-orig',
+    labels = 'inferred',
+    label_mode = 'categorical',
+    color_mode = 'rgb',
+    image_size = (300, 300),
+    shuffle = False,
+    validation_split = 0.05,
+    subset = 'both',
+    verbose = True
+)
+
+from m import mod