update

jzhang-github · Jan 2, 2024 · dc896dd · dc896dd
1 parent 4f7a8a3
commit dc896dd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/agat/__init__.py b/agat/__init__.py
@@ -15,5 +15,6 @@
 from .data.build_dataset import CrystalGraph
 from .model.fit import Fit
 from .app.cata.high_throughput_predict import HtAds
+from .model.model import PotentialModel
 
 del os
diff --git a/agat/model/__init__.py b/agat/model/__init__.py
@@ -7,4 +7,5 @@
 
 from .layer import Layer
 from .model import PotentialModel, CrystalPropertyModel, AtomicPropertyModel, AtomicVectorModel
-from .fit import Fit
+from .fit import Fit
+from agat.lib.model_lib import save_model, load_model, save_state_dict, load_state_dict