minor adjust

stanfordnlp · Mar 26, 2024 · 1d924c1 · 1d924c1
1 parent becd113
commit 1d924c1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pyvene/models/intervenable_base.py b/pyvene/models/intervenable_base.py
@@ -1439,7 +1439,10 @@ def forward(
                 )
 
             # run intervened forward
-            counterfactual_outputs = self.model(**base, labels=labels)
+            if labels is not None:
+                counterfactual_outputs = self.model(**base, labels=labels)
+            else:
+                counterfactual_outputs = self.model(**base)
             set_handlers_to_remove.remove()
 
             self._output_validation()