Fix GPU device and interface errors in models

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

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

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

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()

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):