在TPU上进行训练时出现“不可用：套接字已关闭”错误

Question

在TPU上训练变压器模型时，出现以下错误：

UnavailableError: 2 root error(s) found.
  (0) Unavailable: Socket closed
  (1) Invalid argument: Unable to find a context_id matching the specified one (13089686768223941123). Perhaps the worker was restarted, or the context was GC'd?

我的数据根据​​序列长度分为多个存储桶，以获得最佳性能：

1. 长度小于或等于8

2. 从9到16

3. 从17到24

我正在从随机存储桶中加载每一批。

当我第一次访问每个存储桶时-tensorflow内核会创建一个新图并回溯模型。

在第三次跟踪中发生的错误。因此，如果我从任何两个铲斗中训练，我都不会出错。

Answer 1

据我了解-这是tf 2.3中的错误。

我切换到2.2.0，错误消失了。