fix checks

deepmodeling · Jan 8, 2024 · 0e34c46 · 0e34c46
1 parent 7258e76
commit 0e34c46
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py
@@ -7,3 +7,12 @@
     save_dp_model,
     traverse_model_dict,
 )
+
+__all__ = [
+    "EmbeddingNet",
+    "NativeLayer",
+    "NativeNet",
+    "load_dp_model",
+    "save_dp_model",
+    "traverse_model_dict",
+]
diff --git a/deepmd_utils/model_format/network.py b/deepmd_utils/model_format/network.py
@@ -341,13 +341,14 @@ def __init__(
     ):
         layers = []
         i_in = in_dim
+        rng = np.random.default_rng()
         for idx, ii in enumerate(neuron):
             i_ot = ii
             layers.append(
                 NativeLayer(
-                    np.random.normal(size=(i_in, i_ot)),
-                    b=np.random.normal(size=(ii)),
-                    idt=np.random.normal(size=(ii)) if resnet_dt else None,
+                    rng.normal(size=(i_in, i_ot)),
+                    b=rng.normal(size=(ii)),
+                    idt=rng.normal(size=(ii)) if resnet_dt else None,
                     activation_function=activation_function,
                     resnet=True,
                 ).serialize()
@@ -386,5 +387,5 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
         """
         layers = data.pop("layers")
         obj = cls(**data)
-        super(EmbeddingNet, obj).__init__(layers)
+        super().__init__(layers)
         return obj