diff --git a/torch_dftd_static/nn/dftd3_module.py b/torch_dftd_static/nn/dftd3_module.py
index 3cc7ac9..22a79ba 100644
--- a/torch_dftd_static/nn/dftd3_module.py
+++ b/torch_dftd_static/nn/dftd3_module.py
@@ -7,12 +7,12 @@
 from ase.units import Bohr
 from torch import Tensor
 from torch_dftd.functions.dftd3 import d3_autoang, d3_autoev
-from torch_dftd_static.functions.dftd3 import edisp as edisp_triu
+from torch_dftd_static.functions.dftd3 import edisp
 
 
 """
 Check that c6ab array (shape = (95,95,5,5,3)) has the following structure,
-as assumed in edisp function:
+as assumed in the edisp implementation (not in the original torch-dftd).
 - c6ab[..., 1] is constant along axes 1 (Z of second atom) and 3 (cn of second atom),
   except for Z=0 which does not represent valid atom.
   Second condition in torch.all(...) below exists because there can be
@@ -94,7 +94,7 @@ def calc_energy(
         
         # TODO: add interface for force and stress
         
-        E_disp = d3_autoev * edisp_triu(
+        E_disp = d3_autoev * edisp(
             Z,
             pos = pos / d3_autoang,
             shift_vecs = shift_vecs / d3_autoang,