Update

[ghstack-poisoned]
pytorch · Dec 15, 2024 · b26b9d3 · b26b9d3
1 parent 5259e6e
commit b26b9d3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/_utils.py b/torchrl/_utils.py
@@ -880,11 +880,11 @@ def compile_with_warmup(*args, warmup: int, **kwargs):
         >>> # First 10 calls use the original model
         >>> # After 10 calls, the model is compiled and used
     """
-
     if len(args):
         model = args[0]
+        args = ()
     else:
-        model = kwargs.get("model")
+        model = kwargs.pop("model", None)
     if model is None:
         return lambda model: compile_with_warmup(model, warmup=warmup, **kwargs)
     else: