diff --git a/matgl/models/_m3gnet.py b/matgl/models/_m3gnet.py
index 34a7ad41..494346c8 100644
--- a/matgl/models/_m3gnet.py
+++ b/matgl/models/_m3gnet.py
@@ -49,7 +49,7 @@
 class M3GNet(nn.Module, IOMixIn):
     """The main M3GNet model."""
 
-    __version__ = 1
+    __version__ = 2
 
     def __init__(
         self,
diff --git a/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.json b/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.json
index 9936aa92..fbdd3737 100644
--- a/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.json
+++ b/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.json
@@ -7,7 +7,7 @@
         "model": {
             "@class": "M3GNet",
             "@module": "matgl.models._m3gnet",
-            "@model_version": 1,
+            "@model_version": 2,
             "init_args": {
                 "element_types": [
                     "H",
diff --git a/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.pt b/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.pt
index e1439809..2351011b 100644
Binary files a/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.pt and b/pretrained_models/M3GNet-MP-2018.6.1-Eform/model.pt differ
diff --git a/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.json b/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.json
index 9332650b..1f146d80 100644
--- a/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.json
+++ b/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.json
@@ -7,7 +7,7 @@
         "model": {
             "@class": "M3GNet",
             "@module": "matgl.models._m3gnet",
-            "@model_version": 1,
+            "@model_version": 2,
             "init_args": {
                 "element_types": [
                     "H",
diff --git a/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.pt b/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.pt
index 16c04ca9..cac2754a 100644
Binary files a/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.pt and b/pretrained_models/M3GNet-MP-2021.2.8-DIRECT-PES/model.pt differ
diff --git a/pretrained_models/M3GNet-MP-2021.2.8-PES/model.json b/pretrained_models/M3GNet-MP-2021.2.8-PES/model.json
index 2511d572..4adf89ff 100644
--- a/pretrained_models/M3GNet-MP-2021.2.8-PES/model.json
+++ b/pretrained_models/M3GNet-MP-2021.2.8-PES/model.json
@@ -7,7 +7,7 @@
         "model": {
             "@class": "M3GNet",
             "@module": "matgl.models._m3gnet",
-            "@model_version": 1,
+            "@model_version": 2,
             "init_args": {
                 "element_types": [
                     "H",
diff --git a/pretrained_models/M3GNet-MP-2021.2.8-PES/model.pt b/pretrained_models/M3GNet-MP-2021.2.8-PES/model.pt
index 1f8b8c5d..90396d26 100644
Binary files a/pretrained_models/M3GNet-MP-2021.2.8-PES/model.pt and b/pretrained_models/M3GNet-MP-2021.2.8-PES/model.pt differ