diff --git a/dataset_builders/pie/brat/brat.py b/dataset_builders/pie/brat/brat.py
index 56b55db6..752de772 100644
--- a/dataset_builders/pie/brat/brat.py
+++ b/dataset_builders/pie/brat/brat.py
@@ -4,16 +4,12 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import datasets
-import pytorch_ie
-from pytorch_ie.annotations import (
-    BinaryRelation,
-    LabeledMultiSpan,
-    LabeledSpan,
-    _post_init_single_label,
-)
+from pytorch_ie.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 from pytorch_ie.core import Annotation, AnnotationList, annotation_field
 from pytorch_ie.documents import TextBasedDocument
 
+from pie_datasets import GeneratorBasedBuilder
+
 logger = logging.getLogger(__name__)
 
 
@@ -288,7 +284,7 @@ def __init__(self, merge_fragmented_spans: bool = False, **kwargs):
         self.merge_fragmented_spans = merge_fragmented_spans
 
 
-class BratDatasetLoader(pytorch_ie.data.builder.GeneratorBasedBuilder):
+class BratDatasetLoader(GeneratorBasedBuilder):
     # this requires https://github.com/ChristophAlt/pytorch-ie/pull/288
     DOCUMENT_TYPES = {
         "default": BratDocument,