Thank you!
Could you clarify on the final dense layer having an equal amount of units to the number of classes? The paper I’m implementing isn’t clear - it has 128 layers before the final softmax layer - there aren’t 128 classes. In this basic model with my data, there are only 5 classes. The labels are numbered [1, 5].
I’ve changed to SparseCategoricalCrossEntropy and added a flatten layer, in such a way:
class KnownDetector(Model):
    def __init__(self):
        super(KnownDetector, self).__init__()
        self.TCP = tf.keras.Sequential([
            layers.Conv1D(filters=32, kernel_size=3, activation="relu", input_shape=(100, 2)),  # 100 packets, 2 features
            layers.MaxPool1D(pool_size=3, padding='same'),
            layers.Conv1D(filters=32, kernel_size=3, activation="relu"),  # 100 packets, 2 features
            layers.MaxPool1D(pool_size=3, padding='same'),
            layers.Conv1D(filters=32, kernel_size=3, activation="relu"),  # 100 packets, 2 features
            layers.Dense(128),
            layers.Flatten(),
            tf.keras.layers.Softmax()
        ])
    def call(self, x):
        flag = True  # check if x is TCP or UDP in the future
        if flag:
            return self.TCP(x)
#

fx = KnownDetector()
fx.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(), optimizer='adam', metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
fx.fit(train, train_labels, epochs=1, validation_data=(test, test_labels))

This actually trains the model with ~40% accuracy.
If I try to make the last dense layer into layers.Dense(5), the accuracy is in the single digits. This doesn’t make sense to me - if the final layer should be the number of classes, how does having 128 of them fit at all? I don’t have 128 labels. Could there be an issue with my call function?
Note: changing the SoftMax layer to have 5 dimensions bring a completely different error.