I'm at a complete loss with this one. I have implemented the following code, using the common pattern for defining a 'resolver' and the TPU Strategy.

However, when I run this, the tuner.search function is running at the speed of a regular CPU and I have no idea how to get it to utilise TPUs.

Really appreciate everyone's help!! :)

import keras_tuner as kt
import time

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

def model_builder(hp):
    model = tf.keras.Sequential()

    hp_units = hp.Int('units', min_value=5, max_value=20, step=5)
    model.add(tf.keras.layers.Dense(units=hp_units, input_dim=X.shape[1], activation="relu"))
    model.add(tf.keras.layers.Dense(10))

    hp_learning_rate = hp.Choice('learning_rate', values=[1e-2, 1e-3, 1e-4])

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=hp_learning_rate), 
        loss="binary_crossentropy", 
        metrics=["accuracy"])

    return model

#with strategy.scope():
tuner = kt.Hyperband(
    model_builder, 
    objective="accuracy", 
    max_epochs=10, 
    factor=3, 
    distribution_strategy=strategy)

# Early Stopping Callback
#stop_early = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=5)

tic = time.time()
tuner.search(X, Y, epochs=50, callbacks=[tf.keras.callbacks.EarlyStopping(patience=1)])
toc = time.time()

search_time = toc - tic