I’m trying to train a torch::jit::script::Module but the loss is not decreasing. Here is a minimal example.
#include <torch/script.h>
#include <torch/torch.h>
#include <iostream>
#include <vector>

// custom loader code
#include "nets/nets.h"
#include "util/runfiles.h"

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

  // Custom code that loads the module on CUDA
  auto runfiles = MakeRunfiles(argv[0]);
  torch::jit::script::Module script_module = LoadSegnetBackbone(*runfiles);
  script_module.train();
  std::cout << "Loaded script module" << std::endl;

  // Pull parameters out of the script module so we can push them into the
  // optimizer.
  std::vector<at::Tensor> parameters;
  for (const auto& parameter : script_module.get_parameters()) {
    parameters.push_back(parameter.value().toTensor());
  }
  torch::optim::SGD optimizer(std::move(parameters), /*lr=*/0.01);

  constexpr int kBatchSize = 1;
  for (int epoch = 1; epoch <= 1000; ++epoch) {
    optimizer.zero_grad();

    // The input is a (kBatchSize,3,300,300) tensor filled with ones
    at::Tensor input = torch::ones({kBatchSize, /*channels (rgb) =*/3,
                                    /*height=*/300, /*width=*/300})
                           .to(at::kFloat)
                           .to(at::kCUDA);

    // Push the input through the script module
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(input);
    at::Tensor script_module_forward = script_module.forward(inputs).toTensor();
    // The result is an output tensor of size (kBatchSize, 32, 300, 300)

    // ground truth is a (kBatchSize, 300, 300) tensor filled with ones
    at::Tensor ground_truth =
        torch::ones({kBatchSize, /*height=*/300, /*width=*/300})
            .to(at::kLong)
            .to(at::kCUDA);

    at::Tensor loss = torch::nll_loss2d(
        torch::log_softmax(script_module_forward, /*dim=*/1), ground_truth);
    loss.backward();
    optimizer.step();

    if (epoch % 50 == 0) {
      std::cout << "Loss was " << loss.item<float>() << std::endl;
    }
  }
}

And the output is
Nets example
Loaded script module
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
Loss was 3.44751
...

Am I doing something wrong here?