class_names = train_dataset.classes
class_names