replay messenger

pyro-ppl · Jan 16, 2024 · 1885e22 · 1885e22
1 parent b4564c1
commit 1885e22
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pyro/poutine/replay_messenger.py b/pyro/poutine/replay_messenger.py
@@ -3,11 +3,11 @@
 
 from typing import TYPE_CHECKING, Dict, Optional
 
-import torch
-
 from pyro.poutine.messenger import Messenger
 
 if TYPE_CHECKING:
+    import torch
+
     from pyro.poutine.runtime import Message
     from pyro.poutine.trace_struct import Trace
 
@@ -43,7 +43,7 @@ class ReplayMessenger(Messenger):
     def __init__(
         self,
         trace: Optional["Trace"] = None,
-        params: Optional[Dict[str, torch.Tensor]] = None,
+        params: Optional[Dict[str, "torch.Tensor"]] = None,
     ) -> None:
         """
         :param trace: a trace whose values should be reused