diff --git a/sdata/mappers/sample_mappers.py b/sdata/mappers/sample_mappers.py
index 82a5918..193cf1d 100644
--- a/sdata/mappers/sample_mappers.py
+++ b/sdata/mappers/sample_mappers.py
@@ -152,7 +152,7 @@ def __call__(self, x: Dict) -> Dict:
         if self.skip_this_sample(x):
             return x
         if self.use_data_key:
-            h, w = map(lambda y: x["json"][y], (self.h_key, self.w_key))
+            h, w = map(lambda y: x[self.data_key][y], (self.h_key, self.w_key))
         else:
             h, w = map(lambda y: x[y], (self.h_key, self.w_key))
         x["original_size_as_tuple"] = torch.tensor([h, w])