Skip to content

Commit 6c715f8

Browse files
[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.
1 parent a69ef52 commit 6c715f8

File tree

1 file changed

+1
-1
lines changed
  • server/text_generation_server/layers

1 file changed

+1
-1
lines changed

server/text_generation_server/layers/bnb.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ def forward(self, x: torch.Tensor):
7070
return out
7171

7272

73-
class Linear4bit(nn.Module):
73+
class Linear4bit(torch.nn.Module):
7474
def __init__(self, weight, bias, quant_type):
7575
super().__init__()
7676
self.weight = Params4bit(

0 commit comments

Comments
 (0)