fix forward hook (#1027)

ModelCloud · Jan 5, 2025 · c88f50f · c88f50f
1 parent 55c594b
commit c88f50f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/gptqmodel/nn_modules/hooked_linear.py b/gptqmodel/nn_modules/hooked_linear.py
@@ -19,7 +19,7 @@ def from_conv1d(conv1d: Conv1D):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         output = super().forward(input)
         if self.forward_hook:
-            self.forward_hook(self, input, output)
+            self.forward_hook(self, (input,), output)
         return output
 
 
@@ -41,7 +41,7 @@ def from_conv2d(conv2d: torch.nn.Conv2d):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         output = super().forward(input)
         if self.forward_hook:
-            self.forward_hook(self, input, output)
+            self.forward_hook(self, (input,), output)
         return output
 
 
@@ -64,7 +64,7 @@ def from_linear(linear: torch.nn.Linear):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         output = super().forward(input)
         if self.forward_hook:
-            self.forward_hook(self, input, output)
+            self.forward_hook(self, (input,), output)
         return output