tensorflowのstrategy.experimental_run_v2の使用する際の注意点

tensorflowには、分散学習用のAPIが用意されています。
tensorflowで分散学習する際に必ず使うstrategy.experimental_run_v2は、@tf.functionの中で宣言すべきようです。

    @tf.function
    def distributed_train_step(入力データなど):
            return  mirrored_strategy.experimental_run_v2(学習関数名, args=(入力データなど))

これだけなのですが、mirrored_strategy.experimental_run_v2で使用する学習関数のほうに@tf.funcionを宣言しておくとエラーが発生する場合があるようです。