import torch
import torch.nn as nn
from torch.func import vmap
from torch.export import export

# 1. Inner model (shared across batch)
class TinyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(8, 4)

    def forward(self, x):
        return torch.relu(self.linear(x))

# 2. Module that applies vmap over inner model
class BatchedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = TinyModel()

    def forward(self, x):
        return vmap(self.model)(x)  # vectorize over batch

# 3. Instantiate and test export
x = torch.randn(16, 8)
model = BatchedModel().eval()

# 4. Export
graph_module = export(model, (x,))
print(graph_module.module)

Errors with
in _free_unbacked_symbols_with_path(a, path, real, shape_env, pending, simplify)
   1024 elif isinstance(a, torch.Tensor):
   1025     from torch._subclasses.fake_tensor import FakeTensor
-> 1027     assert isinstance(a, FakeTensor)
   1028     r.update(
   1029         go(
   1030             a.size(),
   (...)
   1033         )
   1034     )
   1035     if a.layout not in [
   1036         torch.sparse_csr,
   1037         torch.sparse_csc,
   1038         torch.sparse_bsr,
   1039         torch.sparse_bsc,
   1040     ]:
AssertionError: