diff --git a/deepmd/pt/__init__.py b/deepmd/pt/__init__.py index ab6173619..daf3d406e 100644 --- a/deepmd/pt/__init__.py +++ b/deepmd/pt/__init__.py @@ -4,6 +4,11 @@ from deepmd.pt.cxx_op import ( ENABLE_CUSTOMIZED_OP, ) +from deepmd.utils.entry_point import ( + load_entry_point, +) + +load_entry_point("deepmd.pt") __all__ = [ "ENABLE_CUSTOMIZED_OP", diff --git a/deepmd/pt/model/__init__.py b/deepmd/pt/model/__init__.py index 8422ac380..6ceb116d8 100644 --- a/deepmd/pt/model/__init__.py +++ b/deepmd/pt/model/__init__.py @@ -1,6 +1 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.utils.entry_point import ( - load_entry_point, -) - -load_entry_point("deepmd.pt")