Fix: return test data and gradient

nipreps · Mar 28, 2024 · 316eb57 · 316eb57
1 parent bcff16d
commit 316eb57
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/eddymotion/data/splitting.py b/src/eddymotion/data/splitting.py
@@ -62,7 +62,9 @@ def lovo_split(dataset, index, with_b0=False):
     mask[index] = True
 
     train_data = data[..., ~mask]
-    train_gradients = gradients[..., ~mask]
+    train_gradients = gradients[..., mask]
+    test_data = data[..., ~mask]
+    test_gradients = gradients[..., mask]
 
     if with_b0:
         train_data = np.concatenate(
@@ -78,5 +80,5 @@ def lovo_split(dataset, index, with_b0=False):
 
     return (
         (train_data, train_gradients),
-        (dwframe, bframe),
+        (test_data, test_gradients),
     )