import torch


def str_dtype_to_torch_dtype(dtype: str) -> torch.dtype:
    """Converts a string representation of a dtype to the corresponding
    PyTorch dtype."""
    if dtype == "int32":
        return torch.int32
    elif dtype == "int64":
        return torch.int64
    elif dtype == "float32":
        return torch.float32
    elif dtype == "float64":
        return torch.float64
    else:
        raise ValueError(f"Unsupported dtype: {dtype}")