fix(torch): segmentation with test_batch_size > 1

jolibrain · Jul 19, 2024 · 0d8d3da · 0d8d3da
1 parent 8bc5b97
commit 0d8d3da
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/backends/torch/torchlib.cc b/src/backends/torch/torchlib.cc
@@ -2242,7 +2242,7 @@ namespace dd
             output = torch::softmax(output, 1);
             torch::Tensor target = batch.target.at(0).to(torch::kFloat64);
             torch::Tensor segmap
-                = torch::flatten(torch::argmax(output.squeeze(), 0))
+                = torch::flatten(torch::argmax(output.squeeze(), 1))
                       .contiguous()
                       .to(torch::kFloat64)
                       .to(cpu); // squeeze removes the batch size