Repository: pytorch
Issue: @custom_op extensions could not be export.export()ed via AOT and run from C++
Labels: module: docs,module: error checking,triaged,module: custom-operators,oncall: pt2,oncall: export,module: pt2-dispatcher

### 🐛 Describe the bug

Here is the repro. I am adding a @custom_op to a working example that saves ExportedProgram via AOT and runs it from C++. When I add custom operation, it stops working : 
Error: Could not find schema for mylib::custom_add.

```
import torch

def custom_add_direct(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a + b

@torch.library.custom_op("mylib::custom_add", mutates_args=(),
                         device_types="cuda",
                         )
def _(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return custom_add_direct(a,b)

@torch.library.register_fake("mylib::custom_add")
def _(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return torch.empty_like(a)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 16)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(16, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        y = self.relu(x)
        x = self.fc2(torch.ops.mylib.custom_add(x, y))
        x = self.sigmoid(x)
        return x

with torch.no_grad():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model().to(device=device)
    example_inputs = (torch.randn(8, 10, device=device),)

    # Export the model                                                                               
    exported = torch.export.export(model, example_inputs)

    # Compile the model                                                                              
    output_path = torch._inductor.aoti_compile_and_package(
        exported,
        example_inputs,
        package_path="model.pt2",
    )

```
Here's the C++ code, it runs model.pt2 perfectly if I replace "torch.ops.mylib.custom_add(x, y)" above with "x+y" : 
``` 
#include <iostream>
#include <vector>
#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_package/model_package_loader.h>

int main() {
    c10::InferenceMode mode;
  
    // Load the compiled model                                                                       
    torch::inductor::AOTIModelPackageLoader loader("model.pt2");
  
    // Prepare input tensor                                                                          
    std::vector<torch::Tensor> inputs = {torch::randn({8, 10}, at::kCUDA)};
  
    // Run inference                                                                                 
    std::vector<torch::Tensor> outputs = loader.run(inputs);
  
    // Print the result                                                                              
    std::cout << "Inference result:" << std::endl;
    std::cout << outputs[0] << std::endl;
  
    return 0;
}
```

### Versions

Pytorch nightly


cc @svekars @brycebortree @sekyondaMeta @AlannaBurke @malfet @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @zou3519 @bdhirsh @yf225