diff --git a/sdv/multi_table/hma.py b/sdv/multi_table/hma.py index a779a347a..12408b01e 100644 --- a/sdv/multi_table/hma.py +++ b/sdv/multi_table/hma.py @@ -635,7 +635,7 @@ def _find_parent_id(likelihoods, num_rows): # Some rows got singular matrix error and the rest were 0 # Fallback to num_rows on the singular matrix rows and # keep 0s on the rest. - likelihoods = likelihoods.fillna(num_rows) + likelihoods = likelihoods.astype(float).fillna(num_rows) else: # at least one row got a valid likelihood, so fill the # rows that got a singular matrix error with the mean diff --git a/tests/integration/multi_table/test_hma.py b/tests/integration/multi_table/test_hma.py index e900f1c43..5cb5c6ba1 100644 --- a/tests/integration/multi_table/test_hma.py +++ b/tests/integration/multi_table/test_hma.py @@ -388,6 +388,7 @@ def test_save_and_load(self, tmp_path): assert loaded_synthesizer.get_info() == synthesizer.get_info() assert loaded_synthesizer.metadata.to_dict() == metadata.to_dict() + @pytest.mark.filterwarnings('error::FutureWarning') def test_hma_primary_key_and_foreign_key_only(self): """Test that ``HMASynthesizer`` can handle tables with primary and foreign keys only.""" # Setup diff --git a/tests/unit/multi_table/test_hma.py b/tests/unit/multi_table/test_hma.py index f95fbc286..a1a776a51 100644 --- a/tests/unit/multi_table/test_hma.py +++ b/tests/unit/multi_table/test_hma.py @@ -1038,6 +1038,7 @@ def test__estimate_num_columns_to_be_modeled_different_distributions(self): assert num_table_cols == estimated_num_columns[table_name] + @pytest.mark.filterwarnings('error::FutureWarning') def test__estimate_num_columns_to_be_modeled(self): """Test the estimated number of columns is exactly the number of columns to be modeled.