drytorch.utils.apply_ops

Module containing functions for nested containers.

Functions

`apply`(obj, expected_type, func)	Extend recursive_apply supports.
`apply_cpu_detach`(obj)	Detach and store in cpu the tensors inside a container.
`apply_to`(obj, device)	Change the device of tensors inside a container.
`recursive_apply`(-> ~drytorch.utils.apply_ops._C)	Look for an expected type and apply a given function.

apply(obj: _C, expected_type: type[_T], func: Callable[[_T], _T]) → _C[source]

Extend recursive_apply supports.

If the input has attributes, it calls recursive_apply, creates a new instance and sets the attributes of a new instance to the new values.

Parameters:

obj (_C) – container or class containing other containers and tensors.
expected_type (type[_T]) – the type of the objects to modify.
func (Callable[[_T], _T]) – a function that modifies objects of the expected type.

Returns:

The container or class with the modified objects.

Return type:

_C

apply_cpu_detach(obj: _C) → _C[source]

Detach and store in cpu the tensors inside a container.

Parameters:: obj (_C) – container or class containing other containers and tensors.
Returns:: the same obj with the tensor on cpu.
Return type:: _C

apply_to(obj: _C, device: device) → _C[source]

Change the device of tensors inside a container.

Parameters:

Returns:

the same container with the tensor on the target device.

Return type:

_C