scvi.module.base.auto_move_data#
- class scvi.module.base.auto_move_data(fn)[source]#
Bases:
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__
orforward
.- Parameters:
fn (Callable) – A nn.Module method for which the arguments should be moved to the device the parameters are on.
- Return type: