From 6008cf114df9362373f493d1632dcc2d3d5b3e74 Mon Sep 17 00:00:00 2001
From: zhanghao <zhouyinzhanghao@gmail.com>
Date: Fri, 7 Jun 2024 12:23:57 +0800
Subject: [PATCH] update lem

---
 dptb/nn/embedding/lem.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/dptb/nn/embedding/lem.py b/dptb/nn/embedding/lem.py
index 718eb63c..82791203 100644
--- a/dptb/nn/embedding/lem.py
+++ b/dptb/nn/embedding/lem.py
@@ -110,11 +110,11 @@ def __init__(
         self.sh = SphericalHarmonics(
             irreps_sh, sh_normalized, sh_normalization
         )
-        self.onehot = OneHotAtomEncoding(num_types=n_atom, set_features=False)
+        self.onehot = OneHotAtomEncoding(num_types=self.n_atom, set_features=False)
 
         self.init_layer = InitLayer(
             idp=self.idp,
-            num_types=n_atom,
+            num_types=self.n_atom,
             n_radial_basis=n_radial_basis,
             r_max=r_max,
             irreps_sh=irreps_sh,
@@ -146,7 +146,7 @@ def __init__(
                 irreps_out = irreps_hidden
 
             self.layers.append(Layer(
-                num_types=n_atom,
+                num_types=self.n_atom,
                 # required params
                 avg_num_neighbors=avg_num_neighbors,
                 irreps_in=irreps_in,