Fix and clean up multihost preprocessing.#146

hertschuh merged 1 commit intokeras-team:mainkeras-team/keras-rs:mainfrom

cantonios:jaxcantonios/keras-rs:jaxCopy head branch name to clipboard