add dependency to fix unit tests

Author: philschmid

setup.py

@@ -55,6 +55,7 @@
 
 # Hugging Face specific dependencies
 extras["transformers"] = ["transformers[sklearn,sentencepiece]>=4.17.0"]
+extras["diffusers"] = ["diffusers==0.23.1"]
 
 # framework specific dependencies
 extras["torch"] = ["torch>=1.8.0", "torchaudio"]
@@ -87,8 +88,7 @@
     "flake8>=3.8.3",
 ]
 
-extras["dev"] = extras["transformers"] + extras["mms"] + extras["torch"] + extras["tensorflow"]
-
+extras["dev"] = extras["transformers"] + extras["mms"] + extras["torch"] + extras["tensorflow"] + extras["diffusers"]
 setup(
     name="sagemaker-huggingface-inference-toolkit",
     version=VERSION,

src/sagemaker_huggingface_inference_toolkit/transformers_utils.py

@@ -284,15 +284,14 @@ def get_pipeline(task: str, device: int, model_dir: Path, **kwargs) -> Pipeline:
         kwargs["tokenizer"] = model_dir
 
     if TRUST_REMOTE_CODE and os.environ.get("HF_MODEL_ID", None) is not None and device == 0:
-        torch_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
         tokenizer = AutoTokenizer.from_pretrained(os.environ["HF_MODEL_ID"])
 
         hf_pipeline = pipeline(
             task=task,
             model=os.environ["HF_MODEL_ID"],
             tokenizer=tokenizer,
             trust_remote_code=TRUST_REMOTE_CODE,
-            model_kwargs={"device_map": "auto", "torch_dtype": torch_dtype},
+            model_kwargs={"device_map": "auto", "torch_dtype": "auto"},
         )
     elif is_diffusers_available() and task == "text-to-image":
         hf_pipeline = get_diffusers_pipeline(task=task, model_dir=model_dir, device=device, **kwargs)