connection.py

# Detect hardware
try:
  tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection
except ValueError: # If TPU not found, resolve to GPU instead
  tpu = None
  gpus = tf.config.experimental.list_logical_devices("GPU")
    
# Select appropriate distribution strategy
if tpu:
	tf.tpu.experimental.initialize_tpu_system(tpu)
	strategy = tf.distribute.experimental.TPUStrategy(tpu, steps_per_run=128)
	print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])  
else:
	strategy = tf.distribute.get_strategy() # Default strategy that works on CPU and single GPU
	print('Running on CPU')
print("Number of accelerators: ", strategy.num_replicas_in_sync)