Update Training a M3GNet Formation Energy Model with PyTorch Lightnin…

…g.ipynb Convert collate_fn into collate_fn_graph for structure-wise property prediction. Signed-off-by: Tsz Wai Ko <[email protected]>
materialsvirtuallab · Jul 26, 2024 · b3ba4e7 · b3ba4e7
1 parent 4c980c6
commit b3ba4e7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/Training a M3GNet Formation Energy Model with PyTorch Lightning.ipynb b/examples/Training a M3GNet Formation Energy Model with PyTorch Lightning.ipynb
@@ -34,7 +34,7 @@
     "from tqdm import tqdm\n",
     "\n",
     "from matgl.ext.pymatgen import Structure2Graph, get_element_list\n",
-    "from matgl.graph.data import MGLDataset, MGLDataLoader, collate_fn\n",
+    "from matgl.graph.data import MGLDataset, MGLDataLoader, collate_fn_graph\n",
     "from matgl.models import M3GNet\n",
     "from matgl.utils.io import RemoteFile\n",
     "from matgl.utils.training import ModelLightningModule\n",
@@ -156,7 +156,7 @@
     "    shuffle=True,\n",
     "    random_state=42,\n",
     ")\n",
-    "my_collate_fn = partial(collate_fn, include_line_graph=True)\n",
+    "my_collate_fn = partial(collate_fn_graph, include_line_graph=True)\n",
     "train_loader, val_loader, test_loader = MGLDataLoader(\n",
     "    train_data=train_data,\n",
     "    val_data=val_data,\n",