small fix

scverse · Oct 30, 2024 · 77df3e2 · 77df3e2
1 parent 2045dff
commit 77df3e2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/scvi/nn/_base_components.py b/src/scvi/nn/_base_components.py
@@ -510,7 +510,7 @@ def forward(
         # The decoder returns values for the parameters of the ZINB distribution
         raw_px_scale = self.factor_regressor(z, cont_covs, *cat_list)
         px_scale = torch.softmax(raw_px_scale, dim=-1)
-        px_dropout = self.px_dropout_decoder(z, *cat_list)
+        px_dropout = self.px_dropout_decoder(z, cont_covs, *cat_list)
         px_rate = torch.exp(library) * px_scale
         px_r = None