Conditionally importing pomegranate

Author: amontanez24

sdmetrics/single_table/bayesian_network.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 import pandas as pd
-import torch
 
 from sdmetrics.goal import Goal
 from sdmetrics.single_table.base import SingleTableMetric
@@ -18,6 +17,7 @@ class BNLikelihoodBase(SingleTableMetric):
     @classmethod
     def _likelihoods(cls, real_data, synthetic_data, metadata=None, structure=None):
         try:
+            import torch
             from pomegranate.bayesian_network import BayesianNetwork
         except ImportError:
             raise ImportError(

tests/unit/single_table/test_bayesian_network.py

@@ -44,7 +44,7 @@ def metadata():
 
 
 class TestBNLikelihood:
-    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None})
+    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None, 'torch': None})
     def test_compute_error(self):
         """Test that an `ImportError` is raised."""
         # Setup
@@ -71,7 +71,7 @@ def test_compute(self, real_data, synthetic_data, metadata):
 
 
 class TestBNLogLikelihood:
-    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None})
+    @patch.dict('sys.modules', {'pomegranate.bayesian_network': None, 'torch': None})
     def test_compute_error(self):
         """Test that an `ImportError` is raised."""
         # Setup