From 208f442d5712004f0364f16cb42ee402b48db29b Mon Sep 17 00:00:00 2001
From: Richie Bendall <richiebendall@gmail.com>
Date: Sun, 27 Oct 2024 13:55:37 +1300
Subject: [PATCH] Fix type hints

---
 hyperlight/convert.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/hyperlight/convert.py b/hyperlight/convert.py
index 1625844..a4609e1 100644
--- a/hyperlight/convert.py
+++ b/hyperlight/convert.py
@@ -85,8 +85,8 @@ def hypernetize_single(
 
 def hypernetize(
     model: nn.Module,
-    modules: Optional[List[nn.Module]] = None,
-    parameters: Optional[List[nn.Parameter]] = None,
+    modules: Optional[Union[List[nn.Module], Dict[str, nn.Module]]] = None,
+    parameters: Optional[Union[List[nn.Parameter], Dict[str, nn.Parameter]]] = None,
     return_values: bool = False,
     inplace: bool = True,
 ):