Note you can use this on Zygote to preallocate stuff: https://github.com/oxinabox/AutoPreallocation.jl https://github.com/oxinabox/AutoPreallocation.jl . It doe

Note you can use this on Zygote to preallocate stuff: https://github.com/oxinabox/AutoPreallocation.jl https://github.com/oxinabox/AutoPreallocation.jl . It doesn't support GPUs yet, mainly because the dev on it needs a GPU CI setup, but it should mostly just work (issue https://github.com/oxinabox/AutoPreallocation.jl/issues/10 https://github.com/oxinabox/AutoPreallocation.jl/issues/10)