How to split operations in each layer so some run on CPU, others run on GPU

[mjkpolo]

Hi,

I was wondering if it's possible to select certain operations to run on the CPU and others for the GPU.
For example, I want to take a model like Llama and try using the CPU for operations with low arithmetic intensity like RMSNorm and see what the impact is.

Thanks!