scvi.module.base.auto_move_data

Contents

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__ 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:

Callable