fix torch incompatibility of numpy

constantinpape · Nov 9, 2022 · 2c1a832 · 2c1a832
1 parent 61339e1
commit 2c1a832
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/torch_em/transform/raw.py b/torch_em/transform/raw.py
@@ -144,10 +144,13 @@ def __init__(self, multiplier=(1.0, 10.0), clip_kwargs={'a_min': 0, 'a_max': 1})
         self.multiplier = multiplier
         self.clip_kwargs = clip_kwargs
 
-    def __call__(self, img, multiplier=None):
+    def __call__(self, img):
         multiplier = np.random.uniform(self.multiplier[0], self.multiplier[1])
         offset = img.min()
-        poisson_noise = np.random.poisson((img - offset) * multiplier) / multiplier + offset
+        poisson_noise = np.random.poisson((img - offset) * multiplier)
+        if isinstance(img, torch.Tensor):
+            poisson_noise = torch.Tensor(poisson_noise)
+        poisson_noise = poisson_noise / multiplier + offset
         if self.clip_kwargs:
             return np.clip(poisson_noise, **self.clip_kwargs)
         return poisson_noise