From 8ab29eeb750ba4970999cc1ec1e1a1f02a6b6741 Mon Sep 17 00:00:00 2001
From: Tyler Osterberg <tyoster@amazon.com>
Date: Fri, 19 Jul 2024 17:32:58 -0700
Subject: [PATCH] [ci] fix hf lora test for accelerate

---
 tests/integration/llm/prepare.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
index 8d4b93900..0e5597fb0 100644
--- a/tests/integration/llm/prepare.py
+++ b/tests/integration/llm/prepare.py
@@ -47,6 +47,7 @@
     "gpt4all-lora": {
         "option.model_id": "s3://djl-llm/gpt4all-lora/",
         "option.tensor_parallel_degree": 4,
+        "option.device_map": "auto",
         "option.task": "text-generation",
         "option.dtype": "fp16"
     }