Set device using parent class constructor

Signed-off-by: Nilaksh Das <[email protected]>

Author: nilakshdas

art/defences/preprocessor/mp3_compression_pytorch.py

@@ -67,19 +67,17 @@ def __init__(
         import torch  # lgtm [py/repeated-import]
         from torch.autograd import Function
 
-        super().__init__(is_fitted=True, apply_fit=apply_fit, apply_predict=apply_predict)
+        super().__init__(
+            device_type=device_type,
+            is_fitted=True,
+            apply_fit=apply_fit,
+            apply_predict=apply_predict,
+        )
         self.channels_first = channels_first
         self.sample_rate = sample_rate
         self.verbose = verbose
         self._check_params()
 
-        # Set device
-        if device_type == "cpu" or not torch.cuda.is_available():
-            self._device = torch.device("cpu")
-        else:  # pragma: no cover
-            cuda_idx = torch.cuda.current_device()
-            self._device = torch.device("cuda:{}".format(cuda_idx))
-
         self.compression_numpy = Mp3Compression(
             sample_rate=sample_rate,
             channels_first=channels_first,

art/defences/preprocessor/spatial_smoothing_pytorch.py

@@ -74,20 +74,17 @@ def __init__(
         """
         import torch  # lgtm [py/repeated-import]
 
-        super().__init__(apply_fit=apply_fit, apply_predict=apply_predict)
+        super().__init__(
+            device_type=device_type,
+            apply_fit=apply_fit,
+            apply_predict=apply_predict,
+        )
 
         self.channels_first = channels_first
         self.window_size = window_size
         self.clip_values = clip_values
         self._check_params()
 
-        # Set device
-        if device_type == "cpu" or not torch.cuda.is_available():
-            self._device = torch.device("cpu")
-        else:  # pragma: no cover
-            cuda_idx = torch.cuda.current_device()
-            self._device = torch.device("cuda:{}".format(cuda_idx))
-
         from kornia.filters import MedianBlur
 
         class MedianBlurCustom(MedianBlur):

art/defences/preprocessor/video_compression_pytorch.py

@@ -73,20 +73,18 @@ def __init__(
         import torch  # lgtm [py/repeated-import]
         from torch.autograd import Function
 
-        super().__init__(is_fitted=True, apply_fit=apply_fit, apply_predict=apply_predict)
+        super().__init__(
+            device_type=device_type,
+            is_fitted=True,
+            apply_fit=apply_fit,
+            apply_predict=apply_predict,
+        )
         self.video_format = video_format
         self.constant_rate_factor = constant_rate_factor
         self.channels_first = channels_first
         self.verbose = verbose
         self._check_params()
 
-        # Set device
-        if device_type == "cpu" or not torch.cuda.is_available():
-            self._device = torch.device("cpu")
-        else:  # pragma: no cover
-            cuda_idx = torch.cuda.current_device()
-            self._device = torch.device("cuda:{}".format(cuda_idx))
-
         self.compression_numpy = VideoCompression(
             video_format=video_format,
             constant_rate_factor=constant_rate_factor,

art/preprocessing/audio/l_filter/pytorch.py

@@ -73,21 +73,19 @@ def __init__(
         """
         import torch  # lgtm [py/repeated-import]
 
-        super().__init__(is_fitted=True, apply_fit=apply_fit, apply_predict=apply_predict)
+        super().__init__(
+            device_type=device_type,
+            is_fitted=True,
+            apply_fit=apply_fit,
+            apply_predict=apply_predict,
+        )
 
         self.numerator_coef = numerator_coef
         self.denominator_coef = denominator_coef
         self.clip_values = clip_values
         self.verbose = verbose
         self._check_params()
 
-        # Set device
-        if device_type == "cpu" or not torch.cuda.is_available():
-            self._device = torch.device("cpu")
-        else:  # pragma: no cover
-            cuda_idx = torch.cuda.current_device()
-            self._device = torch.device("cuda:{}".format(cuda_idx))
-
     def forward(
         self, x: "torch.Tensor", y: Optional["torch.Tensor"] = None
     ) -> Tuple["torch.Tensor", Optional["torch.Tensor"]]:

art/preprocessing/standardisation_mean_std/pytorch.py

@@ -56,7 +56,12 @@ def __init__(
         """
         import torch  # lgtm [py/repeated-import]
 
-        super().__init__(is_fitted=True, apply_fit=apply_fit, apply_predict=apply_predict)
+        super().__init__(
+            device_type=device_type,
+            is_fitted=True,
+            apply_fit=apply_fit,
+            apply_predict=apply_predict,
+        )
         self.mean = np.asarray(mean, dtype=ART_NUMPY_DTYPE)
         self.std = np.asarray(std, dtype=ART_NUMPY_DTYPE)
         self._check_params()
@@ -65,13 +70,6 @@ def __init__(
         self._broadcastable_mean = None
         self._broadcastable_std = None
 
-        # Set device
-        if device_type == "cpu" or not torch.cuda.is_available():
-            self._device = torch.device("cpu")
-        else:  # pragma: no cover
-            cuda_idx = torch.cuda.current_device()
-            self._device = torch.device("cuda:{}".format(cuda_idx))
-
     def forward(
         self, x: "torch.Tensor", y: Optional["torch.Tensor"] = None
     ) -> Tuple["torch.Tensor", Optional["torch.Tensor"]]: