Google Cloud VM по-прежнему использует CPU вместо TPU для выполнения скрипта Python/Tensorflow

Я настроил машину TPU в Google Cloud и думаю, что сделал это правильно, потому что когда я запускаю ctpu status, он возвращает RUNNING.

Однако у меня есть скрипт Python, который я пытаюсь запустить, и я хочу, чтобы он использовал TPU. Однако он все еще использует ЦП, судя по первым нескольким строкам вывода в терминале. Выход

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 7191847218438877393
, name: "/device:XLA_CPU:0"
device_type: "XLA_CPU"
memory_limit: 17179869184
locality {
}
incarnation: 6058516396679200559
physical_device_desc: "device: XLA_CPU device"
]

Команда, которую я запускаю: python3 test.py 1 --tpu-name=$TPU_NAME

Я уже запускал export TPU_NAME=tpu_vm1 и подтверждал это с помощью echo

Так что я могу делать неправильно? Как я мог бы вместо этого использовать скрипт TPU?

На всякий случай вот отредактированный кроме моего test.py скрипт:

#
# resnet time-to-accuracy-improvement tests
#

import os
from numpy.random import seed
seed(1)
import tensorflow as tf
tf.random.set_seed(2)
import numpy
import time

import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50
import mycallbacks

from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.layers import Input
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.utils import plot_model

...

# display device type
from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())