diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py
index 4ca21413..8005b33d 100644
--- a/lerobot/common/policies/diffusion/modeling_diffusion.py
+++ b/lerobot/common/policies/diffusion/modeling_diffusion.py
@@ -28,7 +28,7 @@ from lerobot.common.policies.utils import (
 )
 
 
-def _make_noise_scheduler(name: str, kwargs: dict) -> DDPMScheduler | DDIMScheduler:
+def _make_noise_scheduler(name: str, **kwargs: dict) -> DDPMScheduler | DDIMScheduler:
     """
     Factory for noise scheduler instances of the given name. All kwargs are passed
     to the scheduler.