Script runs on v5p-8 but gets stucked on xmp.spawn on v5p-32

xla
1

Hi, I struggle with running scripts on any accelerator type that has more than 8 TPU cores. I wonder if it is related to the flags like --internal-ips or --tunnel-through-iap

gcloud compute tpus tpu-vm create ${TPU_NAME}
–project ${PROJECT_ID}
–zone ${ZONE}
–accelerator-type ${ACCELERATOR_TYPE}
–version ${RUNTIME_VERSION}
–subnetwork projects/${PROJECT_ID}/regions/${REGION}/subnetworks/${SUBNET_NAME}
–internal-ips

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}
–project ${PROJECT_ID} --zone ${ZONE} --worker=all --tunnel-through-iap
–command='…