AttributeError: модуль «лен» не имеет атрибута «nn»

Question

Я пытаюсь запустить RegNeRF, для чего требуется лен. При установке последней версии flax==0.6.0 я получил сообщение об ошибке, указывающее, что flax не имеет атрибута optim. В этом ответе предлагается понизить версию льна до 0.5.1. При этом теперь я получаю сообщение об ошибкеAttributeError: module 'flax' has no attribute 'nn'

Я не мог найти никаких решений в Интернете для этой ошибки. Любая помощь приветствуется.

Я использую убунту 20.04