import tensorflow as tf

# Check available devices
print(tf.config.list_physical_devices())

# Use GPU for computation
with tf.device('/GPU:0'):
    # TensorFlow operations

# Use TPU for computation
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://ip_address:8470')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():
    # TensorFlow operations

        __
        __