Preserve comment

pyro-ppl · Dec 18, 2023 · 6d4a030 · 6d4a030
1 parent 9151f43
commit 6d4a030
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py
@@ -149,7 +149,11 @@ def _reset(self):
 
     def _process_message(self, msg):
         frame = CondIndepStackFrame(
-            self.name, self.dim, self.subsample_size, self.counter, self.size
+            name=self.name,
+            dim=self.dim,
+            size=self.subsample_size,
+            counter=self.counter,
+            full_size=self.size,  # used for param initialization
         )
         msg["cond_indep_stack"] = (frame,) + msg["cond_indep_stack"]
         if isinstance(self.size, torch.Tensor) or isinstance(