google cloud platform – How to enable GPU in Vertex AI Pipeline

I’m relatively new to Vertex AI Pipeline and I’m facing an issue where I’m unable to enable GPU utilization within my pipeline. Despite specifying the hardware type for a particular component, it still runs on the CPU, as indicated in the figure below. I suspect there might be a missing configuration in my pipeline setup.

embeddings_task = (
        generate_article_embeddings(
            transformer_model=TRANSFORMER_MODEL,
            article_dataset=articles.output,
        )
        .set_cpu_limit("32")
        .set_memory_limit("208G")
        .add_node_selector_constraint("NVIDIA_TESLA_T4")
        .set_gpu_limit("4")
    )

enter image description here

Could someone please provide guidance on how to correctly enable GPU utilization in a Vertex AI Pipeline?

Any help would be greatly appreciated.

I’ve installed all the necessary GPU libraries to enable GPU acceleration, but when I run the code, it still defaults to using the CPU.

Read more here: Source link