Provide better support for running embeddings on GPU

This commit is contained in:
Nicolas Mowen 2025-06-02 08:36:36 -06:00
parent 931fbc2c22
commit 512d5f6190

View File

@ -63,6 +63,7 @@ COPY --from=rocm /opt/rocm-dist/ /
FROM deps-prelim AS rocm-prelim-hsa-override0
ENV HSA_ENABLE_SDMA=0
ENV MIGRAPHX_ENABLE_NHWC=1
ENV TF_ROCM_USE_IMMEDIATE_MODE=1
COPY --from=rocm-dist / /