fix: copy model to cpu for quantized inference

es-ude · Oct 6, 2023 · 0c5d88e · 0c5d88e
1 parent e18f46f
commit 0c5d88e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/elasticai/creator/nn/fixed_point/precomputed/precomputed_module.py b/elasticai/creator/nn/fixed_point/precomputed/precomputed_module.py
@@ -53,5 +53,5 @@ def _stepped_inputs(self, x: torch.Tensor) -> torch.Tensor:
     def _quantized_inference(self, x: int) -> int:
         fxp_input = self._config.as_rational(x)
         with torch.no_grad():
-            output = self(torch.tensor(fxp_input))
+            output = self.cpu()(torch.tensor(fxp_input))
         return self._config.as_integer(float(output.item()))