google-colaboratory - 関数の最後にある torch_xla rendezvous により、「Failed to meet rendezvous」エラーが発生する

Question

現在、Google Colab の複数の tpu コアでコードを実行しようとしていますが、ターゲット関数の最後で同期コード (xm.rendezvous) が呼び出されるとエラーが発生するようですが、同期コードが上。次に例を示します。

# "Map function": acquires a corresponding Cloud TPU core, creates a tensor on it,
# and prints its core
def simple_map_fn(index, flags):
  #   xm.rendezvous('init') # place rendezvous here instead of at the bottom works fine.

  # Acquires the (unique) Cloud TPU core corresponding to this process's index
  device = xm.xla_device()
  ordinal = xm.get_ordinal()
  local_ordinal = xm.get_ordinal()

  print(f"index {index}, process device {device}, local ordinal {local_ordinal}, ordinal {ordinal}")


  # Barrier to prevent master from exiting before workers connect.
  xm.rendezvous('leave')

# Spawns eight of the map functions, one for each of the eight cores on
# the Cloud TPU
flags = {}

xmp.spawn(simple_map_fn, args=(flags,), nprocs=8, start_method='fork')

このノートブックのように Google Colab で上記のコードを実行すると、次のエラーが発生します。

Exception in device=TPU:7: tensorflow/compiler/xla/xla_client/mesh_service.cc:294 : Failed to meet rendezvous 'leave': Socket closed (14)

ランデブーがターゲット関数の下部に配置されたときにランデブーが失敗する理由は何ですか?