You should be getting a batched dataset object, where labels are represented by integers. You cannot use square brackets to extract elements from it. That’s the meaning of “not subscriptable” error.
To see what is inside the first batch you can use a for loop:
for x, y in train_dataset.take(1):
    print('Inputs:', x, 'Labels:', y)

I’m not sure why you want to separate the classes, but it could be accomplished with a filter() function. Please, see the documentation here: tf.data.Dataset | TensorFlow Core v2.6.0 1