From b475541aef14b185fb4fd5e6e0983a2e83ca6eeb Mon Sep 17 00:00:00 2001
From: Yerdos Ordabayev <yordabay@broadinstitute.org>
Date: Mon, 30 Oct 2023 12:57:14 +0000
Subject: [PATCH] fix import

---
 pyro/poutine/runtime.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py
index 2c51f8189f..441b771080 100644
--- a/pyro/poutine/runtime.py
+++ b/pyro/poutine/runtime.py
@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
+from typing_extensions import TypedDict
 
 from pyro.params.param_store import (  # noqa: F401
     _MODULE_NAMESPACE_DIVIDER,
@@ -14,8 +15,6 @@
 )
 
 if TYPE_CHECKING:
-    from typing_extensions import TypedDict
-
     from pyro.poutine.indep_messenger import CondIndepStackFrame
     from pyro.poutine.messenger import Messenger