From a3537160dc2f3ca5d28fb57f9572a71ce0ea3f33 Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Thu, 23 Feb 2023 09:17:40 +0100
Subject: [PATCH] fix autocast issue (#121)

---
 src/peft/tuners/lora.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 0a8e1ca..6756b0c 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -494,8 +494,18 @@ if is_bnb_available():
 
         def forward(self, x: torch.Tensor):
             result = super().forward(x)
+
             if self.disable_adapters:
                 return result
             elif self.r > 0:
-                result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
+                if not torch.is_autocast_enabled():
+                    expected_dtype = result.dtype
+
+                    if x.dtype != torch.float32:
+                        x = x.float()
+                    output = self.lora_B(self.lora_A(self.lora_dropout(x))).to(expected_dtype) * self.scaling
+                    result += output
+                else:
+                    output = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
+                    result += output
             return result