diff --git a/tests/test_hf_model.py b/tests/test_hf_model.py
index ea45fd5..f1d91f4 100644
--- a/tests/test_hf_model.py
+++ b/tests/test_hf_model.py
@@ -1,5 +1,5 @@
 import pytest
-import torch
+import torch, auto_gptq
 from auto_gptq.modeling import BaseGPTQForCausalLM
 from PIL import Image
 from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer