Can you share the exact train, test, val json files that you used . The ones with the dataset seems to be incompatible with the code.