Squeeze fixes for import

tracel-ai · May 19, 2024 · b5bbaf8 · b5bbaf8
1 parent 79793ca
commit b5bbaf8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/crates/burn-import/src/burn/node/squeeze.rs b/crates/burn-import/src/burn/node/squeeze.rs
@@ -24,10 +24,10 @@ impl<PS: PrecisionSettings> NodeCodegen<PS> for SqueezeNode {
         let input = scope.tensor_use_owned(&self.input, node_position);
         let output = &self.output.name;
 
-        let axis = &self.axes.first().unwrap().to_tokens();
+        let axes_arg = &self.axes.to_tokens();
 
         quote! {
-            let #output = #input.squeeze_dims(#axis);
+            let #output = #input.squeeze_dims(&#axes_arg);
         }
     }