explicit MlpRemat(const std::vector<int>& sizes) {
      int forward_tensor = -1;
      int backward_tensor = -1;
      int op = -1;
      // Forward pass:
      for (const int size : sizes) {
        op = AddOperation(/*is_stateful=*/false);
        if (forward_tensor >= 0) AddUse(op, forward_tensor);
        forward_tensor = AddTensor(size);
        AddUse(op, forward_tensor);
      }