From 4831bbbefde87a0f9f033954564584cec83e5743 Mon Sep 17 00:00:00 2001
From: BowenD-UCB <84425382+BowenD-UCB@users.noreply.github.com>
Date: Wed, 7 Feb 2024 14:03:59 -0800
Subject: [PATCH] v0.3.4 release

---
 chgnet/model/dynamics.py |  2 +-
 chgnet/model/model.py    | 16 ++++++++++++----
 pyproject.toml           |  2 +-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/chgnet/model/dynamics.py b/chgnet/model/dynamics.py
index bd27858..97b9743 100644
--- a/chgnet/model/dynamics.py
+++ b/chgnet/model/dynamics.py
@@ -89,7 +89,7 @@ def __init__(
                 self.device = f"cuda:{cuda_devices_sorted_by_free_mem()[-1]}"
 
         # Move the model to the specified device
-        self.model = (model or CHGNet.load()).to(self.device)
+        self.model = (model or CHGNet.load(verbose=False)).to(self.device)
         self.model.graph_converter.set_isolated_atom_response(on_isolated_atoms)
         self.stress_weight = stress_weight
         print(f"CHGNet will run on {self.device}")
diff --git a/chgnet/model/model.py b/chgnet/model/model.py
index e944225..01cc79e 100644
--- a/chgnet/model/model.py
+++ b/chgnet/model/model.py
@@ -669,16 +669,23 @@ def from_file(cls, path, **kwargs) -> CHGNet:
         return CHGNet.from_dict(state["model"], **kwargs)
 
     @classmethod
-    def load(cls, model_name="0.3.0", use_device: str | None = None) -> CHGNet:
+    def load(
+        cls,
+        model_name="0.3.0",
+        use_device: str | None = None,
+        verbose: bool = True,
+    ) -> CHGNet:
         """Load pretrained CHGNet model.
 
         Args:
-            model_name (str, optional): Defaults to "0.3.0".
+            model_name (str, optional):
+                Default = "0.3.0".
             use_device (str, optional): The device to be used for predictions,
                 either "cpu", "cuda", or "mps". If not specified, the default device is
                 automatically selected based on the available options.
                 Default = None
-
+            verbose (bool): whether to print model device information
+                Default = True
         Raises:
             ValueError: On unknown model_name.
         """
@@ -709,7 +716,8 @@ def load(cls, model_name="0.3.0", use_device: str | None = None) -> CHGNet:
 
         # Move the model to the specified device
         model = model.to(device)
-        print(f"CHGNet will run on {device}")
+        if verbose:
+            print(f"CHGNet will run on {device}")
         return model
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 8a77d3b..0319d88 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chgnet"
-version = "0.3.3"
+version = "0.3.4"
 description = "Pretrained Universal Neural Network Potential for Charge-informed Atomistic Modeling"
 authors = [{ name = "Bowen Deng", email = "bowendeng@berkeley.edu" }]
 requires-python = ">=3.9"