Cosine Annealing with Warm Restart LR Scheduler recieving an unexpected kwarg T_i

[turoger]

Describe the bug

Calling CosineAnnealingWarmRestart as an option in pykeen.pipeline.pipeline results in the following error:

CosineAnnealingWarmRestarts.__init__() got an unexpected keyword argument 'T_i

I think this error is because lr_scheduler_kwargs has initialized a T_i state when T_i is not required to run the scheduler (T_i equal to T_0 at initialization by Torch so no need to pass T_i as part of the state)

How to reproduce

res = pipeline(
    dataset=wn.WN18,
    model="RotatE",
    model_kwargs=dict(embedding_dim=200),
    training_kwargs=dict(num_epochs=5, batch_size=128),
    optimizer="Adam",
    optimizer_kwargs={"lr": 0.01},
    lr_scheduler="CosineAnnealingWarmRestarts",
    lr_scheduler_kwargs={
        "T_0": 10,
    },
)

Environment

Key	Value
OS	posix
Platform	Linux
Release	5.15.0-100-generic
Time	Tue Apr 2 21:03:50 2024
Python	3.11.8
PyKEEN	1.10.2
PyKEEN Hash	UNHASHED
PyKEEN Branch
PyTorch	2.2.1+cu118
CUDA Available?	true
CUDA Version	11.8
cuDNN Version	8700

Additional information

No response

Issue Template Checks

• This is not a feature request (use a different issue template if it is)
• This is not a question (use the discussions forum instead)
• I've read the text explaining why including environment information is important and understand if I omit this information that my issue will be dismissed

[mberr]

Hi @turoger ,

the issue should be fixed with #1386 being merged to master. It will be included in our next release, but you can already get access by installing from source.