Skip to content

Commit

Permalink
Cleanup build errors
Browse files Browse the repository at this point in the history
  • Loading branch information
agelas committed Feb 13, 2025
1 parent 6f52bc8 commit 07f64e8
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 5 deletions.
3 changes: 1 addition & 2 deletions crates/burn-import/src/burn/node/split.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ impl<PS: PrecisionSettings> NodeCodegen<PS> for SplitNode {
let input = scope.tensor_use_owned(&self.input, node_position);
let axis = self.config.axis.to_tokens();

let split_tensors = syn::Ident::new("split_tensors", proc_macro2::Span::call_site());
let outputs = self
.outputs
.iter()
Expand Down Expand Up @@ -126,7 +125,7 @@ mod tests {
&self,
tensor1: Tensor<B, 2>,
) -> (Tensor<B, 2>, Tensor<B, 2>) {
let split_tensors = tensor1.split(2, 0);
let mut split_tensors = tensor1.split(2, 0);

let [tensor2, tensor3] = split_tensors.try_into().unwrap();
(tensor2, tensor3)
Expand Down
6 changes: 3 additions & 3 deletions crates/onnx-ir/src/dim_inference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -915,7 +915,7 @@ fn split_update_outputs(node: &mut Node) {
};

let input_dims = match &input_tensor.shape {
Some(shape) => shape.iter().map(|dim| *dim).collect::<Vec<usize>>(),
Some(shape) => shape.to_vec(),
None => panic!("Split: Input tensor shape is not defined"),
};

Expand Down Expand Up @@ -970,8 +970,8 @@ fn split_update_outputs(node: &mut Node) {
let mut sizes = vec![split_size; num_outputs];

// According to ONNX spec, the last chunk will be smaller if not evenly divisible
for i in 0..remainder {
sizes[i] += 1;
for size in sizes.iter_mut().take(remainder) {
*size += 1;
}

sizes
Expand Down

0 comments on commit 07f64e8

Please sign in to comment.