diff --git a/multidms/model.py b/multidms/model.py
index c9d67a0..5b6ee14 100644
--- a/multidms/model.py
+++ b/multidms/model.py
@@ -1282,7 +1282,7 @@ def mut_param_heatmap(
                 continue
             addtl_tooltip_stats.append(f"wildtype_{condition}")
             muts_df[f"wildtype_{condition}"] = muts_df.site.apply(
-                lambda site: site_map.loc[int(site), condition]
+                lambda site: site_map.loc[site, condition]
             )
 
         # melt conditions and stats cols, beta is already "tall"
diff --git a/multidms/model_collection.py b/multidms/model_collection.py
index 12f5132..378f225 100644
--- a/multidms/model_collection.py
+++ b/multidms/model_collection.py
@@ -660,7 +660,7 @@ def mut_param_heatmap(
                 continue
             addtl_tooltip_stats.append(f"wildtype_{condition}")
             muts_df[f"wildtype_{condition}"] = muts_df.site.apply(
-                lambda site: self.site_map_union.loc[int(site), condition]
+                lambda site: self.site_map_union.loc[site, condition]
             )
 
         # melt conditions and stats cols, beta is already "tall"
diff --git a/tests/test_data.py b/tests/test_data.py
index 4205faf..cd1e895 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -383,6 +383,32 @@ def test_fit_models():
     assert list(tall_combined.index.names) == ["scale_coeff_lasso_shift"]
 
 
+def test_ModelCollection_charts():
+    """
+    Test fitting two different models in
+    parallel using multidms.model_collection.fit_models
+    """
+    data = multidms.Data(
+        TEST_FUNC_SCORES,
+        alphabet=multidms.AAS_WITHSTOP,
+        reference="a",
+        assert_site_integrity=False,
+    )
+    params = {
+        "dataset": [data],
+        "iterations_per_step": [2],
+        "scale_coeff_lasso_shift": [0.0, 1e-5],
+    }
+    _, _, fit_models_df = multidms.model_collection.fit_models(
+        params,
+        n_threads=min(os.cpu_count(), 4),
+    )
+    mc = multidms.model_collection.ModelCollection(fit_models_df)
+
+    mc.mut_param_heatmap(query="scale_coeff_lasso_shift == 0.0")
+    mc.shift_sparsity()
+
+
 def test_data_names():
     """
     Test that the default data names are correctly