import torch

from training.models.registry import get_sae_model


def create_sae(
    sae_input_dim: int,
    inner_dim: int,
    device: torch.device,
    weight_dtype: torch.dtype,
    sae_type: str = "sae",
    **kwargs,
) -> torch.nn.Module:
    model_cls = get_sae_model(sae_type)
    autoenc_model = model_cls(sae_input_dim, inner_dim, **kwargs)
    autoenc_model.to(device, dtype=weight_dtype)
    autoenc_model.train()
    return autoenc_model