InvalidArgumentError: logits and labels must be broadcastable: logits_size=[32,2] labels_size=[32,200]

I am training a Reset model on Tiny-Imagenet dataset. When I am training the same model on Cifar-10 or Cifar100 I am not facing any errors. However, when I am using Tiny -Imagenet along with ImagedatGenerator() I face this error. Is it because of the way I feed data into the network?

But I am also doubtful about the way dataset is structure. When I train on the smaller dataset of Tiny Imagenet it works without any error. I am witnessing this error when I scale up the dataset.Dataset structure is as follows.

tiny-Imagenet
|_train
|_class1_images
|_class2_images …
|_val
|_images
|_all_the images

Please find my code.

import numpy as np
import h5py
import matplotlib.pyplot as plt
import copy
from scipy.io import savemat,loadmat
import warnings
warnings.filterwarnings('ignore',category=FutureWarning)

import tensorflow
print('Tensorflow version = ',tensorflow.__version__)
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.callbacks import LearningRateScheduler, History
from tensorflow.keras.preprocessing.image import ImageDataGenerator 
from tensorflow.keras.preprocessing import image

from PIL import Image
from sklearn.utils import shuffle


#from tensorflow.keras import backend as K

from ResNetModel import resnet
from Utils import cutout,LR_WarmRestart,GetDataGen,plot_history
from tensorflow.keras.callbacks import TensorBoard
from datetime import datetime


WhichDataSet = 'Tiny-Imagenet'
init_lr = 0.1
epochs = 2 # 254
batch_size = 32 #125
My_wd=5e-4/2
resnet_width = 10
resnet_depth = 20
Loss = 'categorical_crossentropy'
Optimizer = SGD(lr=init_lr,decay=0.0, momentum=0.9, nesterov=False)
Metrics = ['accuracy']
ModelsPath = 'TrainedModels/Tensorflow.keras/'


#In3
#load and prepare data
train_datagen = ImageDataGenerator(rescale=1./255)
val_datagen = ImageDataGenerator(rescale=1./255)
test_dataset = ImageDataGenerator() #resacle?


train_generator = train_datagen.flow_from_directory('Datasets/tiny-imagenet-200/train/', 
                                    batch_size= batch_size, class_mode='categorical')
val_generator = val_datagen.flow_from_directory('Datasets/tiny-imagenet-200/val/',
                                    batch_size=batch_size, class_mode='categorical')

test_generator = test_dataset.flow_from_directory('Datasets/tiny-imagenet-200/test/', 
                                    batch_size= batch_size, class_mode='categorical')

x_train, y_train = next(train_generator)
x_test, y_test = next(val_generator)

print(x_train.shape)
print(x_test.shape)
print(y_train.shape)
print(y_test.shape)



num_classes = np.unique(y_train).shape[0]
y_true = val_generator.classes


input_shape = x_train.shape[1:]
x_train = x_train.astype('float32')/255.0
x_test = x_test.astype('float32')/255.0
y_train = y_train.astype('float32')
y_test = y_test.astype('float32')
#y_train = tensorflow.keras.utils.to_categorical(y_train,num_classes= num_classes)
#y_test = tensorflow.keras.utils.to_categorical(y_test,num_classes= num_classes)



def catcross_entropy_logits_loss():
    def loss(y_true, y_pred):
        return tensorflow.keras.losses.categorical_crossentropy(y_true, y_pred, from_logits=True)
    return loss


#fdefine a datagen or generating training samples with flip and pad
dataGenerator = GetDataGen(UseCutout)

#define and compile the model
model = resnet(input_shape=input_shape, num_classes=num_classes,wd=My_wd,width=resnet_width)
model.compile(loss=catcross_entropy_logits_loss() ,optimizer = Optimizer, metrics = Metrics)

#print  the model
model.summary()

#define the learnng rate schedule
steps_per_epoch = int(np.floor(train_generator.n // batch_size ))

lr_scheduler = LR_WarmRestart(nbatch=steps_per_epoch,
                              initial_lr=init_lr, min_lr=init_lr*1e-4,
                              epochs_restart = [1.0,3.0, 7.0, 15.0, 31.0, 63.0,127.0,255.0]) 

#define callbacks
history = History()
callbacks = [lr_scheduler,history]




history = model.fit_generator(dataGenerator.flow(x_train, y_train, batch_size=batch_size),
                              validation_data=(x_test, y_test),
                              epochs=epochs,
                              verbose=1,
                              callbacks=callbacks,
                              steps_per_epoch =steps_per_epoch)


#In6
#get final performance

y_pred = model.predict(x_test)
print(y_pred)
print('Test accuracy (%):', 100*sum(np.argmax(y_pred,-1)==np.argmax(y_test,-1))/y_test.shape[0])

output:

Tensorflow version =  2.2.0
Found 100000 images belonging to 200 classes.
Found 10000 images belonging to 1 classes.
Found 10000 images belonging to 1 classes.
(32, 256, 256, 3)
(32, 256, 256, 3)
(32, 200)
(32, 1)

Epoch 1/2
Traceback (most recent call last):
  File "/Users/priyanka/Desktop/Code/Tiny_Imagenet.py", line 163, in <module>
    steps_per_epoch =steps_per_epoch)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 324, in new_func
    return func(*args, **kwargs)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1479, in fit_generator
    initial_epoch=initial_epoch)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 66, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 848, in fit
    tmp_logs = train_function(iterator)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 644, in _call
    return self._stateless_fn(*args, **kwds)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2420, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1665, in _filtered_call
    self.captured_inputs)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1746, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 598, in call
    ctx=ctx)
  File "/anaconda3/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  logits and labels must be broadcastable: logits_size=[32,2] labels_size=[32,200]
         [[node loss/softmax_cross_entropy_with_logits