diff --git a/arkindex/training/api.py b/arkindex/training/api.py
index c77e858fe80606e5ae77a3fc6e7a4609c778f94d..b167ea914fe0a55bee56a28102eebd1ecd0bb681 100644
--- a/arkindex/training/api.py
+++ b/arkindex/training/api.py
@@ -716,6 +716,17 @@ class DatasetUpdate(ACLMixin, RetrieveUpdateDestroyAPIView):
         if obj.state == DatasetState.Complete:
             raise ValidationError(detail="This dataset is in complete state and cannot be modified anymore.")
 
+    def update(self, request, *args, **kwargs):
+        # Do exactly the same thing as what DRF does, but without the automatic prefetch cache removal:
+        # https://github.com/encode/django-rest-framework/blob/2da473c8c8e024e80c13a624782f1da6272812da/rest_framework/mixins.py#L70
+        # This allows `set_elements` to still be returned after the update.
+        partial = kwargs.pop("partial", False)
+        instance = self.get_object()
+        serializer = self.get_serializer(instance, data=request.data, partial=partial)
+        serializer.is_valid(raise_exception=True)
+        self.perform_update(serializer)
+        return Response(serializer.data)
+
     def perform_destroy(self, dataset):
         DatasetElement.objects.filter(set__dataset_id=dataset.id).delete()
         dataset.sets.all().delete()