Is there a way to fix a batch size mismatch between output from final model layer and input?

What I’ve done: I have a complete dataset of 898 labels with a total of ~55,000 images. For purposes of speed, I took 10 of those labels and about ~600 images to test the code below. I’ve tried changing the batchSize, modifying the data function, but to no avail.

Problem: Error: Batch size mismatch: output dense_Dense1 has 10; expected 500 based on input conv2d_Conv2D1_input.

Goal: Either change the final output of dense_Dense1 to have 500, or change the expected input of conv2d_Conv2D1_input to only 10.

Complete Code:

var tf = require('@tensorflow/tfjs');
var tfnode = require('@tensorflow/tfjs-node');
var fs = require(`fs`)

const numberOfClasses = 10;

const imageWidth = 500;
const imageHeight = 800;
const imageChannels = 3;

const batchSize = 3;
const epochsValue = 5;

const createImage = async (fileName) => {
  const imageBuffer = await fs.readFileSync(fileName);
  const image = await tfnode.node.decodeImage(imageBuffer);
  return image;
}

const labelArray = indice => Array.from({length: numberOfClasses}, (_, k) => k === indice ? 1 : 0)

async function* data() {
  for (i = 1; i < numberOfClasses+1; i++) {
    for (x = 10; x < 40; x++) {
      const feature = await createImage(`./images/${i}/${i}-${x}.png`) ;
      const label = tf.tensor1d(labelArray(i))
      yield {xs: feature, ys: label};
    }
  }
}

function onBatchEnd(batch, logs) {
  console.log('Accuracy', logs.acc);
}

const main = async () => {
  const model = tf.sequential();

  model.add(tf.layers.conv2d({
    inputShape: [imageWidth, imageHeight, imageChannels],
    filters: 8,
    kernelSize: 5,
    padding: 'same',
    activation: 'relu'
  }));
  model.add(tf.layers.maxPooling2d({
    poolSize: 2,
    strides: 2
  }));

  model.add(tf.layers.conv2d({
    filters: 16,
    kernelSize: 5,
    padding: 'same',
    activation: 'relu'
  }));
  model.add(tf.layers.maxPooling2d({
    poolSize: 3,
    strides: 3
  }));
  
  model.add(tf.layers.flatten());

  model.add(tf.layers.dense({
    units: numberOfClasses,
    activation: 'softmax'
  }));

  model.compile({
    optimizer: 'sgd',
    loss: 'categoricalCrossentropy',
    metrics: ['accuracy']
  });

  model.summary()

  const ds = tf.data.generator(data);

  model.fitDataset(ds, {
    epochs: 5,
    batchSize: 10,
    callbacks: {onBatchEnd}
  }).then(info => {
    console.log('Final accuracy', info.history.acc);
  });

}
main()