From d4e9e3676b3bc2bce09262afa7d137e19654cc2b Mon Sep 17 00:00:00 2001
From: "Lee, Kin Long Kelvin" <kin.long.kelvin.lee@intel.com>
Date: Fri, 20 Dec 2024 10:20:20 -0800
Subject: [PATCH] fix: mapping site.index instead of just site

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

diff --git a/matsciml/datasets/utils.py b/matsciml/datasets/utils.py
index 50413351..46e68758 100644
--- a/matsciml/datasets/utils.py
+++ b/matsciml/datasets/utils.py
@@ -881,7 +881,7 @@ def _all_sites_have_neighbors(neighbors):
         for src_idx, dst_sites in enumerate(neighbors):
             for site in dst_sites:
                 all_src.append(src_idx)
-                all_dst.append(site)
+                all_dst.append(site.index)
                 all_images.append(site.image)
     if any([len(obj) == 0 for obj in [all_src, all_dst, all_images]]):
         raise ValueError(