Chore/tch 0.18.0 (#2542)

* Update tch version * Fix split ops with empty chunks
tracel-ai · Nov 26, 2024 · a6a5c22 · a6a5c22
1 parent 369e148
commit a6a5c22
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/crates/burn-tch/src/ops/base.rs b/crates/burn-tch/src/ops/base.rs
@@ -421,6 +421,7 @@ impl TchOps {
             .tensor
             .split(split_size as i64, dim as i64)
             .into_iter()
+            .filter(|x| x.numel() > 0)
             .map(TchTensor::new)
             .collect()
     }
@@ -435,6 +436,7 @@ impl TchOps {
             .tensor
             .split_with_sizes(split_sizes_i64, dim as i64)
             .into_iter()
+            .filter(|x| x.numel() > 0)
             .map(TchTensor::new)
             .collect()
     }