Questions about Autoregressive rollout (loop in python)

[helloworld-ren]

Hi,

I tried running the gencast_mini_demo.ipynb in Colaboratory but got a error at the cell block entitled. I don't konw I should how to run it. Any idea on how to get the mini demo running successfully ?

See error:

ValueError: Only interpret mode is supported on CPU backend.

[vksastry]

According to me this is coming from the splash attention which currently only works for TPU. So if you want to use it to run on CPU or GPU override the denoiser_architecture_config.sparse_transformer_config.attention_type to triblockdiag_mha.

[andrewlkd]

Indeed, as in #113 (comment), this is because of running on CPU.

[helloworld-ren]

Thank you very much for your reply.I will try your suggestion to run in TPU environment