Why does the validation loss not go down when training my TensorFlow implementation of AlexNet?

Question

I am trying to implement AlexNet in Tensorflow. Matlab has a built-in AlexNet implementation that I use with the cats and dogs dataset (mledu-datasets/cats_and_dogs_filtered.zip). In Matlab, everything works and both training and evaluation loss goes down as epochs go up. However in my TensorFlow implementation validation loss never goes down and gets stuck around 0.6.

Implementation in Matlab is as follows.

And here is my implementation in TF.

model = keras.models.Sequential([
  keras.layers.experimental.preprocessing.Rescaling(1./255, input_shape=(img_height, img_width, 3)),
  # Conv 1
  keras.layers.Conv2D(filters=96, kernel_size=(11,11), strides=(4,4), activation='relu',
                      #input_shape=(img_height, img_width, 3)
                      ),
  #keras.layers.BatchNormalization(),
  keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
  
  # Conv 2
  keras.layers.Conv2D(filters=128, groups=2, kernel_size=(5,5), strides=(1,1), activation='relu', padding="same"),
  keras.layers.BatchNormalization(),
  keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
  # Conv 3
  keras.layers.Conv2D(filters=384, kernel_size=(3,3), strides=(1,1), activation='relu', padding="same"),
  keras.layers.BatchNormalization(),
  # Conv 4
  keras.layers.Conv2D(filters=192, groups=2, kernel_size=(1,1), strides=(1,1), activation='relu', padding="same"),
  keras.layers.BatchNormalization(),
  # Conv 5
  keras.layers.Conv2D(filters=128, groups=2, kernel_size=(1,1), strides=(1,1), activation='relu', padding="same"),
  keras.layers.BatchNormalization(),
  keras.layers.MaxPool2D(pool_size=(3,3), strides=(2,2)),
  
  # FC
  keras.layers.Flatten(),
  keras.layers.Dense(4096, activation='relu'),
  keras.layers.Dropout(0.5),
  keras.layers.Dense(4096, activation='relu'),
  keras.layers.Dropout(0.5),
  keras.layers.Dense(len(class_names), activation='softmax')
])

model.compile(optimizer=tf.optimizers.SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

history = model.fit(train_ds,
                    validation_data=val_ds,
                    epochs=epochs,
                    batch_size=batch_size,
                    shuffle = True,
                    verbose=0, callbacks=[plot_losses])

I am using the same optimizer and learning rate for both. My question is, is the problem has to do with the group convolutions or something else? Everything else seems to the be same. No matter what I try I can not make val_loss go down.