diff --git a/mpmetrics/metrics.py b/mpmetrics/metrics.py
index 4c58d3c..a4b599b 100644
--- a/mpmetrics/metrics.py
+++ b/mpmetrics/metrics.py
@@ -386,7 +386,7 @@ def __init__(self, mem, thresholds, heap, **kwargs):
         self._created.value = time.time()
 
     def _setstate(self, mem, heap):
-        Struct._setstate(self, mem)
+        Struct._setstate(self, mem, heap)
         self.thresholds = tuple(threshold.value for threshold in self._thresholds)
 
     def observe(self, amount, exemplar=None):
diff --git a/mpmetrics/types.py b/mpmetrics/types.py
index 47d8a0b..0fcfc4b 100644
--- a/mpmetrics/types.py
+++ b/mpmetrics/types.py
@@ -150,6 +150,7 @@ def __init__(self, mem, heap):
         self._heap = heap
 
     def _setstate(self, mem, heap):
+        assert heap is not None
         super()._setstate(mem, heap)
         self._heap = heap
 
diff --git a/tests/metric_test.py b/tests/metric_test.py
index 79e2321..cb9ac2f 100644
--- a/tests/metric_test.py
+++ b/tests/metric_test.py
@@ -349,7 +349,8 @@ def test_time(registry, cls, name):
 @pytest.mark.parametrize('cls', (Counter, Gauge, Summary, Histogram))
 def test_pickle(registry, cls):
     metric = cls('name', 'help', labelnames=('l'), registry=registry)
-    pickle.loads(pickle.dumps(metric.labels('x')))
+    metric.labels('x')
+    assert list(metric.collect()) == list(pickle.loads(pickle.dumps(metric)).collect())
 
 class TestEnum:
     @pytest.fixture