Thanks! Writing my own node worked. I mostly copied the code for CppNode and torch::Function removing the templates and accepting lambdas.
For other that may get here, here’s what we ended up writing.
See:

  
      github.com
  
  
    mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.cpp 4
#include "Function.h"
#include <torch/csrc/autograd/functions/accumulate_grad.h>
#include <torch/torch.h>

namespace torch
{
namespace autograd
{

variable_list LanternFunction::apply(
    variable_list args,
    std::function<variable_list(LanternAutogradContext *, variable_list)> forward,
    std::function<variable_list(LanternAutogradContext *, variable_list)> backward)
{

    std::shared_ptr<LanternNode> node(new LanternNode(), deleteNode);
    node->backward_ = backward;

    const size_t num_inputs = args.size();



  This file has been truncated. show original

  
  
    
    
  
  



  
      github.com
  
  
    mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Function.h 3
#include <iostream>

#define LANTERN_BUILD

#include "lantern/lantern.h"

#include <torch/torch.h>

#include "utils.hpp"

namespace torch
{
namespace autograd
{

// Context to save information during forward that can be accessed in backward
struct LanternAutogradContext
{
    LanternAutogradContext() = default;
    LanternAutogradContext(const LanternAutogradContext &other) = delete;


  This file has been truncated. show original

  
  
    
    
  
  


And here is an example call:

  
      github.com
  
  
    mlverse/lantern/blob/cc84e430d258f7fd75e18ba9145887a492e51a35/src/Autograd.cpp#L126-L141 1

auto res = LanternFunction::apply(
    {x, y},
    [&](LanternAutogradContext *ctx, variable_list args) {
        ctx->save_for_backward(args);
        ctx->saved_data["mul"] = mul;
        return variable_list({args[0] + mul * args[1] + args[0] * args[1]});
    },
    [](LanternAutogradContext *ctx, variable_list grad_output) {
        auto saved = ctx->get_saved_variables();
        int mul = ctx->saved_data["mul"].toInt();
        auto var1 = saved[0];
        auto var2 = saved[1];
        variable_list output = {grad_output[0] + grad_output[0] * var2,
                                grad_output[0] * mul + grad_output[0] * var1};
        return output;
    });