def METHOD_NAME(self, batch_dim_size: int) -> None:
    """
    initialize batch dim size base on the first input batch size
    """
    if self.batch_dim_size != -1 and self.batch_dim_size != batch_dim_size:
        raise RuntimeError(
            f"batch dim size is already initialized! "
            f"Found new batch size: {batch_dim_size} not "
            f"matching existing batch dim size: {self.batch_dim_size}!"
        )
    self.batch_dim_size = batch_dim_size