Not able to install 'pycuda' on HuggingFace container

Hi, I am using HuggingFace SageMaker container for ‘token-classification’ task.
I have fine tuned ‘Bert-base-cased’ model and converted it to onnx format and then to tensorrt engine.
To run inference using tensorrt engine, I want to use PyCuda library.
So I have added PyCuda, tensorrt to requirements.txt file (and custom inference code in
But pycuda installation is failing.
Please see the below code and error.


from sagemaker.huggingface import HuggingFaceModel

hub = {
‘HF_TASK’: ‘token-classification’

huggingface_model = HuggingFaceModel(
model_data = model_artifacts,
transformers_version = ‘4.6.1’,
tensorflow_version = ‘2.4.1’,
py_version = ‘py37’,
env = hub

predictor = huggingface_model.deploy(
initial_instance_count = 1,
instance_type = ‘ml.p3.2xlarge’



But pycuda installation is failing with the below error…

Pycuda installation error:-

In file included from src/cpp/cuda.cpp:4:0:
src/cpp/cuda.hpp:14:10: fatal error: cuda.h: No such file or directory
#include <cuda.h>
compilation terminated.
error: command ‘gcc’ failed with exit status 1

ERROR: Failed building wheel for pycuda

This could be happening due to below two reasons…

  1. PATH variable is not properly set as below
    export PATH=“/usr/local/cuda/bin:$PATH”
    export LD_LIBRARY_PATH=“/usr/local/cuda/lib64:$LD_LIBRARY_PATH”

  2. or CUDA toolkit is not installed

Need help in resolving this issue…