From c7021f0f442680f4b89553b7fa71ae9aa8a544d9 Mon Sep 17 00:00:00 2001
From: PanQiWei <594557445@qq.com>
Date: Thu, 31 Aug 2023 18:39:03 +0800
Subject: [PATCH] fix model type changed after calling .to() method

---
 auto_gptq/modeling/_base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/auto_gptq/modeling/_base.py b/auto_gptq/modeling/_base.py
index e3c3c7c..2c708a7 100644
--- a/auto_gptq/modeling/_base.py
+++ b/auto_gptq/modeling/_base.py
@@ -433,7 +433,8 @@ class BaseGPTQForCausalLM(nn.Module, PushToHubMixin):
             return torch.device(device)
 
     def to(self, device: Union[str, torch.device]):
-        return self.model.to(device)
+        self.model.to(device)
+        return self
 
     def forward(self, *args, **kwargs):
         return self.model(*args, **kwargs)