Conditionally importing torch at top

Author: amontanez24

sdmetrics/single_table/bayesian_network.py

@@ -8,6 +8,11 @@
 from sdmetrics.goal import Goal
 from sdmetrics.single_table.base import SingleTableMetric
 
+try:
+    import torch
+except ModuleNotFoundError:
+    torch = None
+
 LOGGER = logging.getLogger(__name__)
 
 
@@ -17,8 +22,9 @@ class BNLikelihoodBase(SingleTableMetric):
     @classmethod
     def _likelihoods(cls, real_data, synthetic_data, metadata=None, structure=None):
         try:
-            import torch
             from pomegranate.bayesian_network import BayesianNetwork
+            if torch is None:
+                raise ImportError
         except ImportError:
             raise ImportError(
                 'Please install pomegranate with `pip install sdmetrics[pomegranate]`.'

tests/unit/single_table/test_bayesian_network.py

@@ -44,9 +44,22 @@ def metadata():
 
 
 class TestBNLikelihood:
-    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None, 'torch': None})
+    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None})
     def test_compute_error(self):
-        """Test that an `ImportError` is raised."""
+        """Test that an `ImportError` is raised when pomegranate isn't installed."""
+        # Setup
+        metric = BNLikelihood()
+
+        # Run and Assert
+        expected_message = re.escape(
+            'Please install pomegranate with `pip install sdmetrics[pomegranate]`.'
+        )
+        with pytest.raises(ImportError, match=expected_message):
+            metric.compute(Mock(), Mock())
+
+    @patch.dict('sys.modules', {'torch': None})
+    def test_compute_error_torch_is_none(self):
+        """Test that an `ImportError` is raised when torch isn't installed."""
         # Setup
         metric = BNLikelihood()
 
@@ -71,7 +84,7 @@ def test_compute(self, real_data, synthetic_data, metadata):
 
 
 class TestBNLogLikelihood:
-    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None, 'torch': None})
+    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None})
     def test_compute_error(self):
         """Test that an `ImportError` is raised."""
         # Setup