Update train.py

train.py

@@ -169,8 +169,8 @@ def main():
         pin_memory=True,
     )
     input_dataset = HorseZebraDataset(
-        root_horse="input/horses",
-        root_zebra="input/zebras",
+        root_horse="input",
+        root_zebra="input",
         transform=config.transforms,
     )
     input_loader = DataLoader(