Device _make_data_loader for prediction

Help scvi-tools
1

I am extending scVi skeleton for my own package. I can train model without problems, but when I try to do prediction my data is no longer on the same device as the pre-trained module.

I use the following data construction in my prediction, similar to other prediction-like functions I have seen in the scVI package:

        adata = self._validate_anndata(adata)
        scdl = self._make_data_loader(
            adata=adata, indices=indices, batch_size=batch_size
        )
        latent = []
        for tensors in scdl:
            inference_inputs = self.module._get_inference_input(tensors)
            outputs = self.module.inference(**inference_inputs)

My validate_adata is a custom function, while make_data_loader comes from the scVI Base Model.
I can not find a way to specify device in the data loader (to match the one of the model). How should I deal with this?

2

So if I understand correctly, the self.module is on GPU and data on CPU.

Consider adding the following decorator to self.module.inference

as in here

this essentially moves the data to the correct device (based on self.module)

3

Thank you. Adding the decorator to the function that cause the issue has resolved it.