Fix model_function_wrapper: unpack args dict as positional (x, t, **c)
apply_model takes (x, t, **c) not (args_dict). Both wrapper call sites were passing the whole dict as x, leaving t unfilled. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
2790dbb8e7
commit
82c76914a2
+2
-2
@@ -341,7 +341,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
pass
|
||||
if _old is not None:
|
||||
return _old(apply_fn, args)
|
||||
return apply_fn(args)
|
||||
return apply_fn(args["input"], args["timestep"], **args.get("c", {}))
|
||||
|
||||
model.model_options["model_function_wrapper"] = _temporal_wrapper
|
||||
|
||||
@@ -358,7 +358,7 @@ def build_enhancements(model, rating_profile, temporal_style, refinement_key, re
|
||||
if _ew is not None:
|
||||
result = _ew(apply_fn, args)
|
||||
else:
|
||||
result = apply_fn(args)
|
||||
result = apply_fn(args["input"], args["timestep"], **args.get("c", {}))
|
||||
|
||||
# Opportunistically save the capture buffer when it has anchor maps.
|
||||
# We do this every call so the last (cleanest) capture wins.
|
||||
|
||||
Reference in New Issue
Block a user