This toy program crashes when trying to call torch::nll_loss. It’s supposed to simulate a pixel-wise classification loss.
I was hoping that the python documentation https://pytorch.org/docs/stable/nn.html#torch.nn.NLLLoss 15 (the subsection on shape) would work for the c++ function.
Here is the toy program
#include <torch/script.h>
#include <torch/torch.h>
#include <iostream>

int main(int argc, char** argv) {
  std::cout << "NLL Example" << std::endl;

  constexpr int kBatchSize = 1;
  at::Tensor net_output =
      torch::log_softmax(torch::randn({kBatchSize, 2, 300, 300}), /*dim=*/1);
  std::cout << "Log softmax succeeded" << std::endl;

  at::Tensor ground_truth = torch::ones({kBatchSize, 300, 300});
  std::cout << "Ground truth created" << std::endl;

  at::Tensor loss = torch::nll_loss(net_output, ground_truth);
  std::cout << "Done" << std::endl;
} 

The program crashes silently with the output
NLL Example
Log softmax succeeded
Ground truth created