SDL editing pixels directly

I'm writing some code for augmenting images in PyTorch (using a custom op written in C++), i.e. my code operates directly on the RGB buffer. I wanted to visualize what my code is doing and figured I could use SDL to directly render the buffer to a window. However, things don't seem to work as I expected, i.e. only the first frame that I render actually gets shown. The relevant code is the following:

int main() {
  bool quit = false;
  SDL_Event event;

  //
  // Init SDL software renderer.
  //

  SDL_Init(SDL_INIT_VIDEO);
  SDL_Window* window =
      SDL_CreateWindow("Data augmenter", SDL_WINDOWPOS_UNDEFINED,
                       SDL_WINDOWPOS_UNDEFINED, WIDTH, HEIGHT, 0);
  SDL_Surface* screen = SDL_GetWindowSurface(window);

  if (screen->format->format != SDL_PIXELFORMAT_ARGB8888) {
    std::cerr << "Unsupported pixel format: "
              << SDL_GetPixelFormatName(screen->format->format) << std::endl;
  }

  //
  // Render loop.
  //

  Renderer renderer(WIDTH, HEIGHT);

  while (!quit) {
    while (SDL_PollEvent(&event)) {
      if (event.type == SDL_QUIT) {
        quit = true;
      }
    }

    SDL_LockSurface(screen);
    renderer.Render(static_cast<uint32_t*>(screen->pixels));
    SDL_UnlockSurface(screen);
    SDL_UpdateWindowSurface(window);
    std::cout << "Rendered frame!" << std::endl;
  }

  SDL_Quit();
  return 0;
}

Am I using using the API somehow wrong?