model cuda loading

dataset/dataset.py

@@ -11,7 +11,7 @@ class Threshold_noise:
         self.threshold = threshold
 
     def __call__(self, x):
-        return torch.clamp(x, min=self.threshold)
+        return torch.where((x <= self.threshold), torch.tensor(0.))
 
 class Log_normalisation:
     """Log normalisation of intensities"""