tensordict.nn.make_functional

tensordict.nn.make_functional(module: nn.Module, funs_to_decorate: Iterable[str] | None = None, keep_params: bool = False, return_params: bool = True) → TensorDict

Converts a nn.Module to a functional module in-place, and returns its params.

Parameters:

• module (torch.nn.Module) – module that is to be made functional.

• funs_to_decorate (iterable of str, optional) – each string must correspond to a function belonging to module. For nested modules, the torch.nn.Module.forward() method will be decorated. Defaults to "forward".

• keep_params (bool, optional) – if True, the module will keep its parameters. Defaults to False.

• return_params (bool, optional) – if True, the parameters will be collected in a nested tensordict and returned. If False, the module will be made functional but still be stateful.