delete outdated MPS comment

CederGroupHub · Sep 26, 2023 · a6b713b · a6b713b
1 parent d5ba24b
commit a6b713b
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 9 deletions.
diff --git a/chgnet/data/dataset.py b/chgnet/data/dataset.py
@@ -520,7 +520,7 @@ def __init__(
 
         self.keys = []
         self.keys = [
-            (mp_id, graph_id) for mp_id, dic in self.data.items() for graph_id in dic
+            (mp_id, graph_id) for mp_id, dct in self.data.items() for graph_id in dct
         ]
         random.shuffle(self.keys)
         print(f"{len(self.data)} mp_ids, {len(self)} structures imported")

diff --git a/chgnet/model/dynamics.py b/chgnet/model/dynamics.py
@@ -80,7 +80,7 @@ def __init__(
         super().__init__(**kwargs)
 
         # Determine the device to use
-        if use_device == "mps" and torch.backends.mps.is_available():
+        if not use_device and torch.backends.mps.is_available():
             self.device = "mps"
         else:
             self.device = use_device or ("cuda" if torch.cuda.is_available() else "cpu")

diff --git a/chgnet/trainer/trainer.py b/chgnet/trainer/trainer.py
@@ -177,9 +177,6 @@ def __init__(
             self.device = use_device
         elif torch.cuda.is_available():
             self.device = "cuda"
-        # mps is disabled until stable version of torch for mps is released
-        # elif torch.backends.mps.is_available():
-        #     self.device = "mps"
         else:
             self.device = "cpu"
         if self.device == "cuda":

diff --git a/examples/fine_tuning.ipynb b/examples/fine_tuning.ipynb
@@ -61,7 +61,7 @@
     "CHGNet is interfaced to [Pymatgen](https://pymatgen.org/), the training samples (normally coming from different DFTs like VASP),\n",
     "need to be converted to [pymatgen.core.structure](https://pymatgen.org/pymatgen.core.html#module-pymatgen.core.structure).\n",
     "\n",
-    "To convert VASP calculation to pymatgen structures and CHGNet labels, you can use the following [code](https://github.com/CederGroupHub/chgnet/blob/main/chgnet/utils/vasp_utils.py):"
+    "To convert VASP calculation to pymatgen structures and CHGNet labels, you can use the following [code](https://github.com/CederGroupHub/chgnet/blob/main/chgnet/utils/vasp_utils.py):\n"
    ]
   },
   {
@@ -79,7 +79,7 @@
     "\n",
     "# ./my_vasp_calc_dir contains vasprun.xml OSZICAR etc.\n",
     "dataset_dict = parse_vasp_dir(file_root=\"./my_vasp_calc_dir\")\n",
-    "print(dataset_dict.keys())"
+    "print(list(dataset_dict))"
    ]
   },
   {
@@ -93,7 +93,7 @@
     "\n",
     "For super-large training dataset, like MPtrj dataset, we recommend [converting them to CHGNet graphs](https://github.com/CederGroupHub/chgnet/blob/main/examples/make_graphs.py). This will save significant memory and graph computing time.\n",
     "\n",
-    "Below are the example codes to save the structures."
+    "Below are the example codes to save the structures.\n"
    ]
   },
   {
@@ -155,7 +155,7 @@
    "id": "e1611921",
    "metadata": {},
    "source": [
-    "## 1. Prepare Training Data"
+    "## 1. Prepare Training Data\n"
    ]
   },
   {