From db7190d6094430d09df62009d43db6c8ebd48446 Mon Sep 17 00:00:00 2001
From: Dhruv Srikanth <51223342+DhruvSrikanth@users.noreply.github.com>
Date: Wed, 15 May 2024 17:03:00 +0100
Subject: [PATCH] [Bug Fix] Update import in quantization layers from nn to
 torch.nn based on import statements in the file header

---
 server/text_generation_server/layers/bnb.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/text_generation_server/layers/bnb.py b/server/text_generation_server/layers/bnb.py
index d27a33a1..ca39919c 100644
--- a/server/text_generation_server/layers/bnb.py
+++ b/server/text_generation_server/layers/bnb.py
@@ -70,7 +70,7 @@ class Linear8bitLt(torch.nn.Module):
         return out
 
 
-class Linear4bit(nn.Module):
+class Linear4bit(torch.nn.Module):
     def __init__(self, weight, bias, quant_type):
         super().__init__()
         self.weight = Params4bit(