From 4f6c3ce2c7a304707ef21feaf7001b833dc5567f Mon Sep 17 00:00:00 2001
From: cslht <cslht011379@gmail.com>
Date: Thu, 6 Aug 2026 01:50:32 -0700
Subject: [PATCH] Fix GPU device and interface errors in models

---
 LION/models/iterative_unrolled/LPD.py     | 12 ++++++++----
 LION/models/iterative_unrolled/cLPD.py    | 14 +++++++++-----
 LION/models/learned_fbp/DeepFBP.py        |  9 +++------
 LION/models/learned_fbp/FusionFBP.py      |  7 +++----
 LION/models/post_processing/FBPConvNet.py |  8 +++-----
 5 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/LION/models/iterative_unrolled/LPD.py b/LION/models/iterative_unrolled/LPD.py
index 0b61dc9e..999322ae 100644
--- a/LION/models/iterative_unrolled/LPD.py
+++ b/LION/models/iterative_unrolled/LPD.py
@@ -189,13 +189,17 @@ def __init__(
                     ]
                 )
         else:
-            self.lambda_dual = (
+            self.register_buffer(
+                "lambda_dual",
                 torch.ones(self.model_parameters.n_iters)
-                * self.model_parameters.step_size
+                * self.model_parameters.step_size,
+                persistent=False,
             )
-            self.lambda_primal = (
+            self.register_buffer(
+                "lambda_primal",
                 torch.ones(self.model_parameters.n_iters)
-                * self.model_parameters.step_size
+                * self.model_parameters.step_size,
+                persistent=False,
             )
 
     @staticmethod
diff --git a/LION/models/iterative_unrolled/cLPD.py b/LION/models/iterative_unrolled/cLPD.py
index 4b695066..7da1dcf0 100644
--- a/LION/models/iterative_unrolled/cLPD.py
+++ b/LION/models/iterative_unrolled/cLPD.py
@@ -382,13 +382,17 @@ def __init__(
                     ]
                 )
         else:
-            self.lambda_dual = (
+            self.register_buffer(
+                "lambda_dual",
                 torch.ones(self.model_parameters.n_iters)
-                * self.model_parameters.step_size
+                * self.model_parameters.step_size,
+                persistent=False,
             )
-            self.lambda_primal = (
+            self.register_buffer(
+                "lambda_primal",
                 torch.ones(self.model_parameters.n_iters)
-                * self.model_parameters.step_size
+                * self.model_parameters.step_size,
+                persistent=False,
             )
 
     @staticmethod
@@ -413,7 +417,7 @@ def continous_LPD_paper():
         params.data_channels = [7, 32, 32, 32, 5]
         params.reg_channels = [6, 32, 32, 32, 5]
         params.learned_step = True
-        params.step_size = True
+        params.step_size = None
         params.step_positive = False
         params.mode = "ct"
         params.do_secon_order = False
diff --git a/LION/models/learned_fbp/DeepFBP.py b/LION/models/learned_fbp/DeepFBP.py
index d7b22ae4..c628d5f6 100644
--- a/LION/models/learned_fbp/DeepFBP.py
+++ b/LION/models/learned_fbp/DeepFBP.py
@@ -168,12 +168,9 @@ def __init__(self, geometry: ct.Geometry, model_parameters: LIONParameter = None
         self.interpolator_conv = Conv1d(1, 1, kernel_size=3, padding=1, bias=False)
 
         # Tomosipo normalization map (1s projection)
-        sinogram_ones = torch.ones(
-            (1, 1, self.num_angles_, self.num_detectors),
-            device=torch.cuda.current_device(),
-        )
-        self.tomosipo_normalizer = (
-            self.AT(sinogram_ones) + 1e-6
+        sinogram_ones = torch.ones((1, 1, self.num_angles_, self.num_detectors))
+        self.register_buffer(
+            "tomosipo_normalizer", self.AT(sinogram_ones) + 1e-6, persistent=False
         )  # Avoid division by zero in normalization
 
         # Denoising blocks
diff --git a/LION/models/learned_fbp/FusionFBP.py b/LION/models/learned_fbp/FusionFBP.py
index aec6d318..55a4aef4 100644
--- a/LION/models/learned_fbp/FusionFBP.py
+++ b/LION/models/learned_fbp/FusionFBP.py
@@ -307,11 +307,10 @@ def __init__(self, geometry: ct.Geometry, model_parameters: LIONParameter = None
         self.interpolator_conv = Conv1d(1, 1, kernel_size=3, padding=1, bias=False)
 
         # Tomosipo normalization map (1s projection)
-        sinogram_ones = torch.ones(
-            (1, 1, self.num_angles_, self.num_detectors),
-            device=torch.cuda.current_device(),
+        sinogram_ones = torch.ones((1, 1, self.num_angles_, self.num_detectors))
+        self.register_buffer(
+            "tomosipo_normalizer", self.AT(sinogram_ones) + 1e-6, persistent=False
         )
-        self.tomosipo_normalizer = self.AT(sinogram_ones) + 1e-6
 
         # Denoising blocks
         self.denoiser = DenoisingBlock()
diff --git a/LION/models/post_processing/FBPConvNet.py b/LION/models/post_processing/FBPConvNet.py
index dae37414..be687a94 100644
--- a/LION/models/post_processing/FBPConvNet.py
+++ b/LION/models/post_processing/FBPConvNet.py
@@ -94,15 +94,13 @@ def forward(self, x):
 class FBPConvNet(LIONmodel):
     def __init__(
         self,
-        geometry_parameters: ct.Geometry,
+        geometry: ct.Geometry,
         model_parameters: Optional[LIONModelParameter] = None,
     ):
 
-        assert (
-            geometry_parameters is not None
-        ), "Geometry parameters required for FBPConvNet."
+        assert geometry is not None, "Geometry parameters required for FBPConvNet."
 
-        super().__init__(model_parameters, geometry_parameters)
+        super().__init__(model_parameters, geometry)
         self._make_operator()
         # standard FBPConvNet (As per paper):