output with shape [] doesn't match the broadcast shape [1] error in Tensor::backward() call with LibTorch backend

[femshima]

I am attemting to implement dice loss with burn. However, the following code leads to an error inside libtorch as in the title of this discussion. Also, when I remove dice.backward();, the error disappears.

This may be an issue in my code, but I have not been able to find the cause of this error.

use burn::backend::Autodiff;
use burn::nn::conv::Conv2dConfig;
use burn::nn::PaddingConfig2d;
use burn::tensor::backend::AutodiffBackend;
use burn::tensor::Tensor;

fn main() {
    type MyBackend = burn::backend::LibTorch;
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::libtorch::LibTorchDevice::Cpu;

    test::<MyAutodiffBackend>(&device);
}

fn test<B: AutodiffBackend>(device: &B::Device) {
    let conv2d = Conv2dConfig::new([16, 1], [3, 3])
        .with_padding(PaddingConfig2d::Same)
        .init(device);

    let x: Tensor<B, 4> = Tensor::zeros([1, 16, 512, 512], device);
    let x: Tensor<B, 4> = conv2d.forward(x);

    let x: Tensor<B, 1> = x.flatten(0, 3);

    let targets: Tensor<B, 4> = Tensor::zeros([1, 1, 512, 512], device);
    let targets: Tensor<B, 1> = targets.flatten(0, 3);

    let intersection: Tensor<B, 1> = (x.clone() * targets.clone()).sum();
    let dice: Tensor<B, 1> = (intersection * 2.0 + 1) / ((x.clone()).sum() + (targets.clone()).sum() + 1);

    dice.backward();
}

Here is the message emitted on panic.

thread 'main' panicked at /home/user/.cargo/registry/src/index.crates.io-6f17d22bba15001f/burn-tch-0.13.2/src/ops/base.rs:295:40:
called `Result::unwrap()` on an `Err` value: Torch("output with shape [] doesn't match the broadcast shape [1]\nException raised from mark_resize_outputs at ../aten/src/ATen/TensorIterator.cpp:1185 (most recent call first):\nframe #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x6c (0x7fd657b9f28c in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libc10.so)\nframe #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xfa (0x7fd657b48f5c in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libc10.so)\nframe #2: at::TensorIteratorBase::mark_resize_outputs(at::TensorIteratorConfig const&) + 0x21a (0x7fd641cd591a in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #3: at::TensorIteratorBase::build(at::TensorIteratorConfig&) + 0x78 (0x7fd641cd59c8 in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #4: at::TensorIteratorBase::build_borrowing_binary_op(at::TensorBase const&, at::TensorBase const&, at::TensorBase const&) + 0xab (0x7fd641cd6a4b in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #5: <unknown function> + 0x2ddafbe (0x7fd6431dafbe in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #6: <unknown function> + 0x50a13dd (0x7fd6454a13dd in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #7: <unknown function> + 0x47da410 (0x7fd644bda410 in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #8: at::_ops::mul__Tensor::call(at::Tensor&, at::Tensor const&) + 0x14a (0x7fd64298493a in /home/user/dev/project/target/release/build/torch-sys-a0156870d86154e2/out/libtorch/libtorch/lib/libtorch_cpu.so)\nframe #9: <unknown function> + 0x472e8 (0x564ebb2962e8 in target/release/project)\nframe #10: <unknown function> + 0x39062 (0x564ebb288062 in target/release/project)\nframe #11: <unknown function> + 0x1d8d1 (0x564ebb26c8d1 in target/release/project)\nframe #12: <unknown function> + 0x2b247 (0x564ebb27a247 in target/release/project)\nframe #13: <unknown function> + 0x22ebd (0x564ebb271ebd in target/release/project)\nframe #14: <unknown function> + 0x4c4f3 (0x564ebb29b4f3 in target/release/project)\nframe #15: <unknown function> + 0x2fcf1 (0x564ebb27ecf1 in target/release/project)\nframe #16: <unknown function> + 0x3794a (0x564ebb28694a in target/release/project)\nframe #17: <unknown function> + 0x368e3 (0x564ebb2858e3 in target/release/project)\nframe #18: <unknown function> + 0x2d0a9 (0x564ebb27c0a9 in target/release/project)\nframe #19: <unknown function> + 0x71711 (0x564ebb2c0711 in target/release/project)\nframe #20: <unknown function> + 0x37a45 (0x564ebb286a45 in target/release/project)\nframe #21: <unknown function> + 0x25d4a (0x7fd657f63d4a in /usr/lib/libc.so.6)\nframe #22: __libc_start_main + 0x8c (0x7fd657f63e0c in /usr/lib/libc.so.6)\nframe #23: <unknown function> + 0x1b585 (0x564ebb26a585 in target/release/project)\n")

[nathanielsimard]

Did you try with another backend?

[femshima]

Thank you for your suggestion.
I tried with Wgpu and NdArray backend, but they did not emit any error (not reproducible with those backends).

[femshima]

Changing

    let intersection: Tensor<B, 1> = (x.clone() * targets.clone()).sum();
    let dice: Tensor<B, 1> = (intersection * 2.0 + 1) / ((x.clone()).sum() + (targets.clone()).sum() + 1);

to

    let intersection: Tensor<B, 1> = (x.clone() * targets.clone()).sum();
    let dice: Tensor<B, 1> = (intersection * 2.0 + 1).reshape([1])
        / ((x.clone()).sum() + (targets.clone()).sum() + 1).reshape([1]);

fixed the issue.

It seems that sum() returns 0-dimension tensor, which results in a dimension mismatch in
https://github.com/tracel-ai/burn/blob/v0.13.2/crates/burn-autodiff/src/ops/tensor.rs#L393.

Maybe related to #1689.

As I found the workaround, I will close discussion.