diff --git a/custom/custom_dataset.py b/custom/custom_dataset.py
index cd02ac6..7f344b5 100644
--- a/custom/custom_dataset.py
+++ b/custom/custom_dataset.py
@@ -103,8 +103,9 @@ def __init__(self, train_file: str=None, ):
         print(self.image[0])
         print("-----------------"*3)
     def __getitem__(self, index):
-        prompt = self.prompt[0]
-        image = Image.open(self.image[0]).convert("RGB")
+        index = index % len(self.image)
+        prompt = self.prompt[index]
+        image = Image.open(self.image[index]).convert("RGB")
         image = self.transform(image)
         
         return image,prompt