RESOLVI RuntimeError: dtype mismatch when dispersion set to "gene-batch"

[tsvvas]

Hi team,

I'm running into a PyTorch dtype mismatch error when training the RESOLVI model. The traceback suggests a matrix multiplication is being attempted between a LongTensor and a FloatTensor.

params = {
        "n_hidden": 32,
        "dropout_rate": 0.05,
        "dispersion": "gene-batch",
        "downsample_counts": False,
    }

scvi.external.RESOLVI.setup_anndata(
    adata,
    layer=LAYER_KEY,
    batch_key=BATCH_KEY,
    prepare_data_kwargs={"spatial_rep": "spatial"},
)
model = scvi.external.RESOLVI(adata, **params)
model.to_device(device)

model.train(**train_params)

The final and most important part of the error is:

File /opt/conda/envs/spatial/lib/python3.11/site-packages/scvi/external/resolvi/_module.py:870, in RESOLVAEGuide.forward(self, x, ind_x, library, y, batch_index, cat_covs, x_n, distances_n, n_obs, kl_weight)
    862 px_r_mle = pyro.param(
    863     "px_r_mle",
    864     self.px_r,
    865     constraint=constraints.greater_than(self.eps),
    866     event_dim=len(self.px_r.shape),
    867 )
    869 if self.dispersion == "gene-batch":
--> 870     px_r_inv = F.linear(
    871         torch.nn.functional.one_hot(batch_index.flatten(), self.n_batch), px_r_mle
    872     )
    873 elif self.dispersion == "gene":
    874     px_r_inv = px_r_mle

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: long int != float

Proposed solution:

The issue seems to be solved by casting the dtype of the result of one_hot() to float:

if self.dispersion == "gene-batch":
    one_hot = torch.nn.functional.one_hot(batch_index.flatten(), self.n_batch).to(px_r_mle.dtype)
    px_r_inv = F.linear(one_hot, px_r_mle)

Versions:

scvi-tools: 1.3.0
PyTorch: 2.6.0+cu124
CUDA: NVIDIA A100

Best regards,
Vasily

[canergen]

Yes, this is correct and fix will be pushed latter.