From 9a88c5dc36b0e04a13d0f32d12d5c61379d5d1fa Mon Sep 17 00:00:00 2001
From: Alexander Kensert <alexander.kensert@gmail.com>
Date: Wed, 10 Apr 2024 15:07:32 +0200
Subject: [PATCH] fix bug

---
 molexpress/datasets/encoders.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/molexpress/datasets/encoders.py b/molexpress/datasets/encoders.py
index 9114622..cb49071 100644
--- a/molexpress/datasets/encoders.py
+++ b/molexpress/datasets/encoders.py
@@ -122,7 +122,7 @@ def __call__(self, molecule: types.Molecule) -> np.ndarray:
 
         if molecule.GetNumBonds() == 0:
             edge_state = np.zeros(
-                shape=(0, self.output_dim), 
+                shape=(0, self.output_dim + int(self.self_loops)), 
                 dtype=self.output_dtype
             )
             return {