Are you running this on GPU or CPU? The performance degradation is expected on GPU as we need the heterogenous runtime support to avoid redundant memory copy between CPU and GPU. @zhiics is currently working on this.
Besides, @jroesch is working on the memory planning for dynamic shape cases to reduce the total number of memory allocations and reuse the buffer as much as possible.