def METHOD_NAME(
    self, module: FusedEmbeddingBagCollection
) -> Dict[str, nn.Parameter]:
    params = {
        name.split(".")[-2]: param
        for name, param in module.state_dict().items()
        if name.endswith(".weight")
    }
    return params