Better documentation for M3GNet potential training with stresses

examples/Training a M3GNet Potential with PyTorch Lightning.ipynb

@@ -135,7 +135,8 @@
     "    shuffle=True,\n",
     "    random_state=42,\n",
     ")\n",
-    "my_collate_fn = partial(collate_fn_pes, include_line_graph=True)\n",
+    "# if you are not intended to use stress for training, switch include_stress=False!\n",
+    "my_collate_fn = partial(collate_fn_pes, include_line_graph=True, include_stress=True)\n",
     "train_loader, val_loader, test_loader = MGLDataLoader(\n",
     "    train_data=train_data,\n",
     "    val_data=val_data,\n",
@@ -148,7 +149,8 @@
     "    element_types=element_types,\n",
     "    is_intensive=False,\n",
     ")\n",
-    "lit_module = PotentialLightningModule(model=model, include_line_graph=True)"
+    "# if you are not intended to use stress for training, set stress_weight=0.0!\n",
+    "lit_module = PotentialLightningModule(model=model, include_line_graph=True, stress_weight=0.01)"
    ]
   },
   {
@@ -383,7 +385,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,

tests/layers/test_core_and_embedding.py

@@ -48,7 +48,7 @@ def test_gated_mlp(self, x):
 
     @pytest.mark.parametrize("normalization", ["layer", "graph"])
     def test_mlp_norm(self, x, graph, normalization):
-        layer = MLP_norm(dims=[10, 3], normalization=normalization, normalize_hidden=True)
+        layer = MLP_norm(dims=[10, 10, 3], normalization=normalization, normalize_hidden=True)
         out = layer(x, g=graph).double()
         assert [out.size()[0], out.size()[1]] == [4, 3]
         assert out.mean().item() == pytest.approx(0, abs=1e-6)