Bug fix about device of subsample in subsample_messenger.py (#3195)

* Bug description about device of subsample in subsample_messenger.py * Bug fix about device of subsample in subsample_messenger.py by Jinnan Huang * Specify the `device` of `self._indices` to fix the device error in subsample_messenger.py * Reformat the line that starts with `new_value = value.index_select(...)`
pyro-ppl · Apr 20, 2023 · 1311b9e · 1311b9e
1 parent 71eeded
commit 1311b9e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pyro/poutine/subsample_messenger.py b/pyro/poutine/subsample_messenger.py
@@ -186,7 +186,9 @@ def _postprocess_message(self, msg):
                     # Subsample parameters with known batch semantics.
                     if self.subsample_size < self.size:
                         value = msg["value"]
-                        new_value = value.index_select(dim, self._indices)
+                        new_value = value.index_select(
+                            dim, self._indices.to(value.device)
+                        )
                         if msg["type"] == "param":
                             if hasattr(value, "_pyro_unconstrained_param"):
                                 param = value._pyro_unconstrained_param