diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index cb56565..34a08c3 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -34,8 +34,8 @@ class aiuNNDataset(torch.utils.data.Dataset):
     
     def __getitem__(self, idx):
         # Get the byte strings
-        low_res_bytes = self.df.iloc[idx]['low_ress']
-        high_res_bytes = self.df.iloc[idx]['high_ress']
+        low_res_bytes = self.df.iloc[idx]['image_512']
+        high_res_bytes = self.df.iloc[idx]['image_1024']
         
         # Convert bytes to numpy arrays
         low_res = cv2.imdecode(np.frombuffer(low_res_bytes, np.uint8), -1)