[Bug Fix] Update torch import reference in bnb quantization (#1902)

# What does this PR do? Fixes `Import Error` occurring from mismatch of usage between torch.nn.Module and nn.Module.
huggingface · May 15, 2024 · 6c715f8 · 6c715f8
1 parent a69ef52
commit 6c715f8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/text_generation_server/layers/bnb.py b/server/text_generation_server/layers/bnb.py
@@ -70,7 +70,7 @@ def forward(self, x: torch.Tensor):
         return out
 
 
-class Linear4bit(nn.Module):
+class Linear4bit(torch.nn.Module):
     def __init__(self, weight, bias, quant_type):
         super().__init__()
         self.weight = Params4bit(