use torch.tensor

Author: R-Palazzo

sdmetrics/single_table/bayesian_network.py

@@ -49,9 +49,9 @@ def _likelihoods(cls, real_data, synthetic_data, metadata=None, structure=None):
         probabilities = []
         for _, row in synthetic_data[fields].iterrows():
             try:
-                probabilities.append(bn.probability([row.to_numpy()]))
+                probabilities.append(torch.tensor(bn.probability([row.to_numpy()])))
             except ValueError:
-                probabilities.append(0)
+                probabilities.append(torch.tensor(0))
 
         return np.asarray(probabilities)
 
@@ -125,7 +125,7 @@ def compute(cls, real_data, synthetic_data, metadata=None, structure=None):
             float:
                 Mean of the probabilities returned by the Bayesian Network.
         """
-        return np.mean(cls._likelihoods(real_data, synthetic_data, metadata, structure))
+        return np.mean(cls._likelihoods(real_data, synthetic_data, metadata, structure)).item()
 
 
 class BNLogLikelihood(BNLikelihoodBase):
@@ -199,7 +199,7 @@ def compute(cls, real_data, synthetic_data, metadata=None, structure=None):
         """
         likelihoods = cls._likelihoods(real_data, synthetic_data, metadata, structure)
         likelihoods[np.where(likelihoods == 0)] = 1e-8
-        return np.mean(np.log(likelihoods))
+        return np.mean(np.log(likelihoods)).item()
 
     @classmethod
     def normalize(cls, raw_score):

tests/unit/single_table/test_bayesian_network.py

@@ -51,7 +51,7 @@ def test_compute(self, real_data, synthetic_data, metadata):
         result = metric.compute(real_data, synthetic_data, metadata)
 
         # Assert
-        assert result == 0.111111104
+        assert result == 0.1111111044883728
 
 
 class TestBNLogLikelihood:
@@ -65,4 +65,4 @@ def test_compute(self, real_data, synthetic_data, metadata):
         result = metric.compute(real_data, synthetic_data, metadata)
 
         # Assert
-        assert result == -7.3347335
+        assert result == -7.334733486175537