minor update

tracel-ai · Sep 28, 2024 · 07f79ab · 07f79ab
1 parent 249d60c
commit 07f79ab
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/crates/burn-import/onnx-tests/tests/top_k/top_k.py b/crates/burn-import/onnx-tests/tests/top_k/top_k.py
@@ -67,7 +67,7 @@ def create_model(op_set_version: int):
 
 def main():
     # Unconmment when initializers are supported.
-    for op_set_version in [1]: #, 10, 11]:
+    for op_set_version in [1, 10, 11]:
         create_model(op_set_version)
 
 

diff --git a/crates/burn-import/src/onnx/op_configuration.rs b/crates/burn-import/src/onnx/op_configuration.rs
@@ -12,7 +12,8 @@ use onnx_ir::ir::{ArgType, AttributeValue, Data, ElementType, Node};
 
 /// Extract and convert a given attribute to i64
 fn extract_attr_value_i64(node: &Node, key: &str) -> i64 {
-    let value = node.attrs.get(key).expect("Expected the following attribute key: {:?}").clone().into_i64();
+    let error_msg = format!("Expected the following attribute key: {:?}", key);
+    let value = node.attrs.get(key).expect(&error_msg).clone().into_i64();
     value
 }