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?