fix: fix embedding net mistake

sbi-dev · Sep 16, 2024 · bcc75db · bcc75db
1 parent 045bf5e
commit bcc75db
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sbi/neural_nets/estimators/categorical_net.py b/sbi/neural_nets/estimators/categorical_net.py
@@ -61,8 +61,8 @@ def __init__(
             self._initialize()
 
     def forward(self, inputs, context=None):
-        embedded_inputs = self.embedding_net.forward(inputs)
-        return super().forward(embedded_inputs, context=context)
+        embedded_context = self.embedding_net.forward(context)
+        return super().forward(inputs, context=embedded_context)
 
     def compute_probs(self, outputs):
         ps = F.softmax(outputs, dim=-1) * self.mask