Repository: pytorch
Issue: How to export onnx with fixed shape output ?
Labels: module: onnx,triaged

### 🐛 Describe the bug

```
import torch

class TRT_SCA(torch.autograd.Function):  
    @staticmethod
    def forward(ctx,
                query,
                key,
                value,
                reference_points,
                spatial_shapes,
                reference_points_cam,
                bev_mask,
                level_start_index):
        out = torch.randn(1, 1600, 256, dtype=torch.float32) 
        return out  # I just want to assign the out shape is [1, 1600, 256]   

    @staticmethod
    def symbolic(g, 
                query,
                key,
                value,
                reference_points,
                spatial_shapes,
                reference_points_cam,
                bev_mask,
                level_start_index):
        return g.op("TRT::SCATT",
                query,
                key,
                value,
                reference_points,
                spatial_shapes,
                reference_points_cam,
                bev_mask,
                level_start_index)
    
trt_sca = TRT_SCA.apply

class SpatialCrossAttention(torch.nn.Module):
    def __init__(self):
        super(SpatialCrossAttention, self).__init__()
    
    def forward(self,
                query,
                key,
                value,
                reference_points=None,
                spatial_shapes=None,
                reference_points_cam=None,
                bev_mask=None,
                level_start_index=None):       
        return trt_sca(
            query,
            key,
            value,
            reference_points,
            spatial_shapes,
            reference_points_cam,
            bev_mask,
            level_start_index)  

query= torch.randn(1, 1600, 256, dtype=torch.float32) 
key= torch.randn(6, 5315, 1, 256, dtype=torch.float32) 
value= torch.randn(6, 5315, 1, 256, dtype=torch.float32) 

reference_points = torch.randn(1, 4, 1600, 3, dtype=torch.float32) 
spatial_shapes= torch.tensor(  [[ 40, 100],
                                [ 20,  50],
                                [ 10,  25],
                                [  5,  13]], dtype=torch.int64)

reference_points_cam=torch.randn(6, 1, 1600, 4, 2, dtype=torch.float32) 
bev_mask=torch.where(torch.randn(6, 1, 1600, 4)  > 0.2, 1, 0)
level_start_index= torch.tensor([ 0, 4000, 5000, 5250], dtype=torch.int64)

nn_model = SpatialCrossAttention()  

print("------------------------------------")

output_file = 'sca.onnx' 
torch.onnx.export(
    nn_model,
    (query,
    key,
    value,
    reference_points,
    spatial_shapes,
    reference_points_cam,
    bev_mask,
    level_start_index),
    output_file,
    export_params=True,
    keep_initializers_as_inputs=True,
    do_constant_folding=True,
    enable_onnx_checker=True, 
    verbose=True,
    opset_version=11,
)
print("export done")
```


### Versions

onnx                      1.15.0     
onnx-graphsurgeon         0.3.21   
onnx-simplifier           0.4.36   
onnxruntime               1.17.1     
torch                     1.10.0+cu113   
torchaudio                0.10.0+cu113   
torchvision               0.11.0+cu113       

### Result   

![image](https://github.com/pytorch/pytorch/assets/38753233/4bda011f-b520-4576-9907-9dd1b73573cf)