Fix ModuleDict/ParameterDict in combination with `torch.jit.scrip…

…t` (#9424) Fixes #9401
pyg-team · Jun 14, 2024 · 1f5e201 · 1f5e201
1 parent f0fa9bf
commit 1f5e201
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/torch_geometric/nn/module_dict.py b/torch_geometric/nn/module_dict.py
@@ -1,4 +1,4 @@
-from typing import Final, Iterable, Mapping, Optional, Set, Tuple, Union
+from typing import Final, Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Module
@@ -11,7 +11,7 @@
 # internal representation and converts it back to `.` in the external
 # representation. It also allows passing tuples as keys.
 class ModuleDict(torch.nn.ModuleDict):
-    CLASS_ATTRS: Final[Set[str]] = set(dir(torch.nn.ModuleDict))
+    CLASS_ATTRS: Final[Tuple[str, ...]] = tuple(dir(torch.nn.ModuleDict))
 
     def __init__(
         self,

diff --git a/torch_geometric/nn/parameter_dict.py b/torch_geometric/nn/parameter_dict.py
@@ -1,4 +1,4 @@
-from typing import Final, Iterable, Mapping, Optional, Set, Tuple, Union
+from typing import Final, Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Parameter
@@ -11,7 +11,7 @@
 # internal representation and converts it back to `.` in the external
 # representation. It also allows passing tuples as keys.
 class ParameterDict(torch.nn.ParameterDict):
-    CLASS_ATTRS: Final[Set[str]] = set(dir(torch.nn.ParameterDict))
+    CLASS_ATTRS: Final[Tuple[str, ...]] = set(dir(torch.nn.ParameterDict))
 
     def __init__(
         self,