v0.3.4 release

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}")

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
 
 

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 = "[email protected]" }]
 requires-python = ">=3.9"