python - 射线错误：缺少 1 个必需的位置参数：“循环”

Question

我正在执行 https://github.com/kingoflolz/mesh-transformer-jax 中给出的 train.py这 使用 Jax、Ray、TPU 来训练这个 mdoel。但是我它在两者之间停止并产生以下错误：

(pid=9454, ip=10.164.0.9) jax runtime initialization starting
2021-07-09 12:11:59,794 ERROR worker.py:78 -- Unhandled error (suppress with RAY_IGNORE_UNHANDLED_ERRORS=1): ray::NetworkRunner.run() (pid=9454, ip=10.164.0.9)
File "python/ray/_raylet.pyx", line 501, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 451, in ray._raylet.execute_task.function_executor
File "/home/paramjeetsingh80/.local/lib/python3.8/site-packages/ray/_private/function_manager.py", line 563, in actor_method_executor
return method(__ray_actor, *args, **kwargs)
File "/home/paramjeetsingh80/mesh-transformer-jax/mesh_transformer/train_actor.py", line 24, in run
TypeError: new() missing 1 required positional argument: 'loops'