From 038fec60b361f2f1cd0cef870a64211753417c48 Mon Sep 17 00:00:00 2001
From: Pierre-Louis Barbarant <plbarbarant@gmail.com>
Date: Wed, 22 May 2024 18:20:49 +0200
Subject: [PATCH] Add nan val check for barycenter features

---
 src/fugw/mappings/sparse_barycenter.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/fugw/mappings/sparse_barycenter.py b/src/fugw/mappings/sparse_barycenter.py
index b75f058..caf3464 100644
--- a/src/fugw/mappings/sparse_barycenter.py
+++ b/src/fugw/mappings/sparse_barycenter.py
@@ -52,6 +52,10 @@ def update_barycenter_features(plans, weights_list, features_list, device):
                 else:
                     barycenter_features += acc
 
+        # Check for NaN values in the barycenter features
+        if torch.isnan(barycenter_features).any():
+            raise ValueError("Barycenter features contain NaN values")
+
         return barycenter_features.T
 
     @staticmethod