Fix (graph/gpxq): fix orig_forward assignment (Xilinx#845)

i-colbert · Feb 13, 2024 · cf37912 · cf37912
1 parent 86304bc
commit cf37912
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/brevitas/graph/gpxq.py b/src/brevitas/graph/gpxq.py
@@ -65,11 +65,10 @@ def __init__(
         self.group_of_parallel_layers = group_of_parallel_layers
         self.return_forward_output = return_forward_output
 
+        self.orig_forward = self.model.forward
         if isinstance(self.model, (GraphModule, TorchGraphModule)):
-            self.orig_forward = self.model.__class__.forward
             self.model.__class__.forward = self.catch_stopfwd
         else:
-            self.orig_forward = self.model.forward
             self.model.forward = self.catch_stopfwd
 
     def _is_module_supported(self, module):