Splitting train/test dataset in Python based comparison operators

I'm working on splitting a training/test dataset based on a comparison operator:

Split the training/testing dataset into two sets: one with class labels < 5 and one with class labels >= 5. Print out the shapes of the resulting two sets from both training and testing datasets.

import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from keras.datasets import mnist
from keras.utils import to_categorical
from tensorflow import keras

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.cifar10.load_data()

I believe the class labels are the train_labels and test_labels. I've usually done this via sci-kit-learn but now I'm being asked to do it via a comparison operator. I didn't even think that was possible. Was wondering if anyone had any ideas on here.

Thanks!