diff --git a/data_structures/src/staking/stakes.rs b/data_structures/src/staking/stakes.rs
index d5da2200d..c65a135c1 100644
--- a/data_structures/src/staking/stakes.rs
+++ b/data_structures/src/staking/stakes.rs
@@ -198,8 +198,8 @@ where
         capability: Capability,
         epoch: Epoch,
         strategy: CensusStrategy,
-    ) -> Box<dyn Iterator<Item = Address> + '_> {
-        let iterator = self.rank(capability, epoch).map(|(address, _)| address);
+    ) -> Box<dyn Iterator<Item = StakeKey<Address>> + '_> {
+        let iterator = self.by_rank(capability, epoch).map(|(address, _)| address);
 
         match strategy {
             CensusStrategy::All => Box::new(iterator),