It looks like you’re trying to use a very old API for data collators, which was deprecated in v3 and removed in v4. DefaultDataCollator() has not collate_batch method, you call it directly on the features you want to batch, and you should implement the __call__ method in your custom data collator.