diff --git a/c/tskit/trees.c b/c/tskit/trees.c
index e8caeb9bd0..dc08f7c3c7 100644
--- a/c/tskit/trees.c
+++ b/c/tskit/trees.c
@@ -6516,7 +6516,7 @@ tsk_treeseq_divergence_matrix_branch(const tsk_treeseq_t *self,
     double tu, tv, d, span, left, right, span_left, span_right;
     double *restrict D;
     sv_tables_t sv;
-    tsk_size_t *ss_offsets = malloc((num_sample_sets + 1) * sizeof(*ss_offsets));
+    tsk_size_t *ss_offsets = tsk_malloc((num_sample_sets + 1) * sizeof(*ss_offsets));
 
     memset(&sv, 0, sizeof(sv));
     ret = tsk_tree_init(&tree, self, 0);
diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py
index 31ad272a05..f86a0ff8bd 100644
--- a/python/tests/test_lowlevel.py
+++ b/python/tests/test_lowlevel.py
@@ -1535,10 +1535,13 @@ def test_divergence_matrix(self):
         windows = [0, ts.get_sequence_length()]
         D = ts.divergence_matrix(windows)
         assert D.shape == (1, n, n)
-        D = ts.divergence_matrix(windows, samples=[0, 1])
+        D = ts.divergence_matrix(windows, sample_set_sizes=[1, 1], sample_sets=[0, 1])
         assert D.shape == (1, 2, 2)
-        D = ts.divergence_matrix(windows, samples=[0, 1], span_normalise=True)
+        D = ts.divergence_matrix(
+            windows, sample_set_sizes=[1, 1], sample_sets=[0, 1], span_normalise=True
+        )
         assert D.shape == (1, 2, 2)
+        # TODO Add moer tests for sample sets handling
         with pytest.raises(TypeError, match="str"):
             ts.divergence_matrix(windows, span_normalise="xdf")
         with pytest.raises(TypeError):
@@ -1547,8 +1550,6 @@ def test_divergence_matrix(self):
             ts.divergence_matrix(windows=[0])
         with pytest.raises(_tskit.LibraryError, match="BAD_WINDOWS"):
             ts.divergence_matrix(windows=[-1, 0, 1])
-        with pytest.raises(ValueError):
-            ts.divergence_matrix(windows=[0, 1], samples="sdf")
         with pytest.raises(ValueError, match="Unrecognised stats mode"):
             ts.divergence_matrix(windows=[0, 1], mode="sdf")
         with pytest.raises(_tskit.LibraryError, match="UNSUPPORTED_STAT_MODE"):