From 452764fa3850003fca12e7d5943e056f18512dd2 Mon Sep 17 00:00:00 2001
From: "Ariel H. Curiale" <curiale@gmail.com>
Date: Wed, 20 Dec 2023 09:05:10 -0500
Subject: [PATCH] GradICON requires the itk_wrapper

---
 src/icon_registration/__init__.py | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/src/icon_registration/__init__.py b/src/icon_registration/__init__.py
index b948879..09df436 100644
--- a/src/icon_registration/__init__.py
+++ b/src/icon_registration/__init__.py
@@ -1,16 +1,7 @@
-from icon_registration.losses import (
-    LNCC,
-    LNCCOnlyInterpolated,
-    BlurredSSD,
-    GradientICON,
-    InverseConsistentNet,
-    gaussian_blur,
-    ssd_only_interpolated,
-    ssd,
-    SSDOnlyInterpolated,
-    SSD,
-    NCC
-)
+from icon_registration.losses import (LNCC, LNCCOnlyInterpolated, BlurredSSD,
+                                      GradientICON, InverseConsistentNet,
+                                      gaussian_blur, ssd_only_interpolated,
+                                      ssd, SSDOnlyInterpolated, SSD, NCC)
 from icon_registration.network_wrappers import (
     DownsampleRegistration,
     FunctionFromMatrix,
@@ -19,3 +10,5 @@
     TwoStepRegistration,
 )
 from icon_registration.train import train_batchfunction, train_datasets
+
+import icon_registration.itk_wrapper