Fix for the pyo3-onnx bindings...

huggingface · Apr 1, 2024 · 0546127 · 0546127
1 parent ad9b80e
commit 0546127
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/candle-pyo3/src/lib.rs b/candle-pyo3/src/lib.rs
@@ -1591,7 +1591,7 @@ fn candle_functional_m(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
 }
 
 #[cfg(feature = "onnx")]
-fn candle_onnx_m(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+fn candle_onnx_m(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     use onnx::{PyONNXModel, PyONNXTensorDescriptor};
     m.add_class::<PyONNXModel>()?;
     m.add_class::<PyONNXTensorDescriptor>()?;
@@ -1608,9 +1608,9 @@ fn candle(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_submodule(&nn)?;
     #[cfg(feature = "onnx")]
     {
-        let onnx = PyModule::new(py, "onnx")?;
-        candle_onnx_m(py, onnx)?;
-        m.add_submodule(onnx)?;
+        let onnx = PyModule::new_bound(py, "onnx")?;
+        candle_onnx_m(py, &onnx)?;
+        m.add_submodule(&onnx)?;
     }
     m.add_class::<PyTensor>()?;
     m.add_class::<PyQTensor>()?;

diff --git a/candle-pyo3/src/onnx.rs b/candle-pyo3/src/onnx.rs
@@ -39,7 +39,7 @@ impl PyONNXTensorDescriptor {
     /// The shape of the tensor.
     /// &RETURNS&: Tuple[Union[int,str,Any]]
     fn shape(&self, py: Python) -> PyResult<Py<PyTuple>> {
-        let shape = PyList::empty(py);
+        let shape = PyList::empty_bound(py);
         if let Some(d) = &self.0.shape {
             for dim in d.dim.iter() {
                 if let Some(value) = &dim.value {