Dimension changing

import beartype
import torch

from jaxtyping import Float, jaxtyped
from torch import Tensor


@beartype.beartype
class Identity(torch.nn.Module):
    @jaxtyped(typechecker=beartype.beartype)
    def forward(self, x: Float[Tensor, "batch d_model"]):
        return x


@jaxtyped(typechecker=beartype.beartype)
def main():
    model = Identity()

    x = torch.ones((4, 1), dtype=torch.float)
    model(x)

    x = torch.ones((6, 1), dtype=torch.float)
    model(x)  # <- Throws an error about "batch" not matching.


if __name__ == "__main__":
    main()

Not sure why this is happening. But this second call to model.forward will complain because "batch" has changed between calls. How can I fix this?

Reproduce the provided Python example using the Identity.forward entry point, jaxtyped, beartype, and PyTorch tensors. Read the shape-checking behavior around the named dimensions and determine why changing the batch size across calls raises an error. Done means the intended batch-size change is either supported or clearly documented with a working fix or explanation.