Yeah, it works:
>>> tf.test.gpu_device_name()
2021-06-17 03:28:08.178927: W tensorflow/core/platform/profile_utils/cpu_utils.cc:108] Failed to find bogomips or clock in /proc/cpuinfo; cannot determine CPU frequency
2021-06-17 03:28:08.181116: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x12a0f740 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2021-06-17 03:28:08.181383: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): Host, Default Version
2021-06-17 03:28:08.219136: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1
2021-06-17 03:28:08.370055: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:08.371122: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x12da3590 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2021-06-17 03:28:08.371250: I tensorflow/compiler/xla/service/service.cc:176] StreamExecutor device (0): NVIDIA Tegra X2, Compute Capability 6.2
2021-06-17 03:28:08.372054: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:08.372345: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1742] Found device 0 with properties:
pciBusID: 0000:00:00.0 name: NVIDIA Tegra X2 computeCapability: 6.2
coreClock: 1.3GHz coreCount: 2 deviceMemorySize: 7.66GiB deviceMemoryBandwidth: 38.74GiB/s
2021-06-17 03:28:08.372531: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.10.2
2021-06-17 03:28:08.382578: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublas.so.10
2021-06-17 03:28:08.396214: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcufft.so.10
2021-06-17 03:28:08.413786: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcurand.so.10
2021-06-17 03:28:08.424710: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusolver.so.10
2021-06-17 03:28:08.433412: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusparse.so.10
2021-06-17 03:28:08.434862: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudnn.so.8
2021-06-17 03:28:08.435363: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:08.436060: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:08.436224: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1884] Adding visible gpu devices: 0
2021-06-17 03:28:08.445436: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.10.2
2021-06-17 03:28:11.674224: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1283] Device interconnect StreamExecutor with strength 1 edge matrix:
2021-06-17 03:28:11.674349: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1289] 0
2021-06-17 03:28:11.674385: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1302] 0: N
2021-06-17 03:28:11.681245: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:11.681632: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1046] ARM64 does not support NUMA - returning NUMA node zero
2021-06-17 03:28:11.682036: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1428] Created TensorFlow device (/device:GPU:0 with 100 MB memory) -> physical GPU (device: 0, name: NVIDIA Tegra X2, pci bus id: 0000:00:00.0, compute capability: 6.2)
'/device:GPU:0'
>>>
Thanks!