Fix type hint for field.

materialsvirtuallab · Jul 29, 2023 · 10fd101 · 10fd101
1 parent b1d2a9f
commit 10fd101
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/matgl/models/_m3gnet.py b/matgl/models/_m3gnet.py
@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import dgl
 import torch
@@ -74,7 +74,7 @@ def __init__(
         use_phi: bool = False,
         niters_set2set: int = 3,
         nlayers_set2set: int = 3,
-        field: str = "node_feat",
+        field: Literal["node_feat", "edge_feat"] = "node_feat",
         include_state: bool = False,
         activation_type: str = "swish",
         **kwargs,