improve code

Author: R-Palazzo

deepecho/models/basic_gan.py

@@ -186,18 +186,15 @@ def __init__(
         if enable_gpu:
             if sys.platform == 'darwin':  # macOS
                 if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
-                    device = torch.device('mps')
+                    device = 'mps'
                 else:
-                    device = torch.device('cpu')
+                    device = 'cpu'
             else:  # Linux/Windows
-                if torch.cuda.is_available():
-                    device = torch.device('cuda')
-                else:
-                    device = torch.device('cpu')
+                device = 'cuda' if torch.cuda.is_available() else 'cpu'
         else:
-            device = torch.device('cpu')
+            device = 'cpu'
 
-        self._device = device
+        self._device = torch.device(device)
         self._verbose = verbose
 
         LOGGER.info('%s instance created', self)

tests/integration/test_basic_gan.py

@@ -30,7 +30,7 @@ def test_deprecation_warning(self):
     def test__init___enable_gpu(self):
         """Test when `enable_gpu` parameter in the constructor."""
         # Setup and Run
-        model = BasicGANModel(epochs=10, enable_gpu=True)
+        model = BasicGANModel(epochs=10)
 
         # Assert
         if (