def wrapper(obj: DataLoader, *args: Any, **kwargs: Any) -> None:
# We need to inspect `init`, as inspecting `obj.__init__`
# can lead to inspecting the wrong function with multiple inheritance
params = inspect.signature(init).parameters
param_names = tuple(param.name for param in params.values()
if param.name != "self" and param.kind not in (
param.VAR_POSITIONAL, param.VAR_KEYWORD))
param_names = param_names[:len(args)]
if not hasattr(obj, "__pl_dl_args"):
obj.__pl_dl_args = args
obj.__pl_dl_kwargs = kwargs
obj.__pl_dl_arg_names = param_names
# We want to use the latest possible value for dataset argument (i.e. ideally what gets passed to DataLoader)
# so that we can be sure, that it will not get changed anymore.
# That is why we are setting this in every `__init__`
if "dataset" in param_names:
setattr(obj, "__dataset", args[param_names.index("dataset")])
elif "dataset" in kwargs:
setattr(obj, "__dataset", kwargs["dataset"])
init(obj, *args, **kwargs)
