fix

factories.py

@@ -50,7 +50,7 @@ class PhysicsWithGenerator(torch.nn.Module):
 
         if self.name == "MotionBlur_medium":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.05),
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.05).to(device_str),
                                              padding="valid", device=device_str)
             self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=0.6, sigma=0.5, device=device_str)
             self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.05, sigma_max=0.05, device=device_str)
@@ -60,7 +60,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                                   "psf_size": 31, "motion_gen_l": 0.6, "motion_gen_s": 0.5}}
         elif self.name == "MotionBlur_hard":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.1),
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=.1).to(device_str),
                                              padding="valid", device=device_str)
             self.physics_generator = MotionBlurGenerator((psf_size, psf_size), l=1.2, sigma=1.0, device=device_str)
             self.generator = self.physics_generator + SigmaGenerator(sigma_min=0.1, sigma_max=0.1, device=device_str)
@@ -70,7 +70,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                                   "psf_size": 31, "motion_gen_l": 1.2, "motion_gen_s": 1.0}}
         elif self.name == "GaussianBlur_easy":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.01),
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.01).to(device_str),
                                              padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=1.0, sigma_max=1.0,
@@ -83,7 +83,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                                   "blur_sigma": 1.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "GaussianBlur_medium":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05),
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05).to(device_str),
                                              padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=2.0, sigma_max=2.0,
@@ -96,7 +96,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                                   "blur_sigma": 2.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "GaussianBlur_hard":
             psf_size = 31
-            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05),
+            self.physics = dinv.physics.Blur(noise_model=dinv.physics.GaussianNoise(sigma=0.05).to(device_str),
                                              padding="valid", device=device_str)
             self.physics_generator = GaussianBlurGenerator(psf_size=(psf_size, psf_size),
                                                            sigma_min=4.0, sigma_max=4.0,
@@ -108,7 +108,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                                  "fixed_params": {"noise_sigma_min": 0.05, "noise_sigma_max": 0.05,
                                                   "blur_sigma": 4.0, "psf_size": 31, "num_channels": 1}}
         elif self.name == "MRI":
-            self.physics = dinv.physics.MRI(noise_model=dinv.physics.GaussianNoise(sigma=.01),
+            self.physics = dinv.physics.MRI(noise_model=dinv.physics.GaussianNoise(sigma=.01).to(device_str),
                                             img_size=(640, 320), device=device_str)
             self.physics_generator = dinv.physics.generator.RandomMaskGenerator((2, 640, 320), acceleration_factor=4)
             self.generator = self.physics_generator
@@ -125,7 +125,7 @@ class PhysicsWithGenerator(torch.nn.Module):
                 circle=False,
                 normalize=True,
                 device=device_str,
-                noise_model=dinv.physics.GaussianNoise(sigma=1e-4),
+                noise_model=dinv.physics.GaussianNoise(sigma=1e-4).to(device_str),
                 max_iter=10,
             )
             self.physics_generator = SigmaGenerator(sigma_min=1e-4, sigma_max=1e-4, device=device_str)
@@ -151,7 +151,6 @@ class PhysicsWithGenerator(torch.nn.Module):
         """Update value of an existing key in save_params."""
         if value != "" and key in list(self.saved_params["updatable_params"].keys()):
             if type(value) == str:  # it may be only a str representation
-                # type: str -> ???
                 value = self.saved_params["updatable_params_converter"][key](value)
             elif isinstance(value, torch.Tensor):
                 value = value.item()  # type: torch.Tensor -> float