scvi.module.base.auto_move_data

class scvi.module.base.auto_move_data(fn)[source]

Decorator for Module methods to move data to correct device.

Input arguments are moved automatically to the correct device. It has no effect if applied to a method of an object that is not an instance of Module and is typically applied to __call__ or forward.

Parameters
fn : Callable

A nn.Module method for which the arguments should be moved to the device the parameters are on.

Return type

CallableCallable