fix kl weight check for function based modules

scverse · Oct 26, 2021 · d8bd388 · d8bd388
1 parent 75ea35d
commit d8bd388
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/scvi/train/_trainingplans.py b/scvi/train/_trainingplans.py
@@ -4,6 +4,7 @@
 import pyro
 import pytorch_lightning as pl
 import torch
+from pyro.nn import PyroModule
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from scvi import _CONSTANTS
@@ -656,9 +657,13 @@ def __init__(
         self.pyro_guide = self.module.guide
         self.pyro_model = self.module.model
 
-        self.use_kl_weight = (
-            "kl_weight" in signature(self.pyro_model.forward).parameters
-        )
+        self.use_kl_weight = False
+        if isinstance(self.pyro_model, PyroModule):
+            self.use_kl_weight = (
+                "kl_weight" in signature(self.pyro_model.forward).parameters
+            )
+        elif callable(self.pyro_model):
+            self.use_kl_weight = "kl_weight" in signature(self.pyro_model).parameters
 
         self.svi = pyro.infer.SVI(
             model=self.pyro_model,