diff --git a/api/controllers/service_api/dataset/segment.py b/api/controllers/service_api/dataset/segment.py
index d4a6b6aa4f..3f7cfcaea8 100644
--- a/api/controllers/service_api/dataset/segment.py
+++ b/api/controllers/service_api/dataset/segment.py
@@ -200,8 +200,8 @@ class DatasetSegmentApi(DatasetApiResource):
         parser.add_argument('segments', type=dict, required=False, nullable=True, location='json')
         args = parser.parse_args()
 
-        SegmentService.segment_create_args_validate(args['segments'], document)
-        segment = SegmentService.update_segment(args['segments'], segment, document, dataset)
+        SegmentService.segment_create_args_validate(args, document)
+        segment = SegmentService.update_segment(args, segment, document, dataset)
         return {
             'data': marshal(segment, segment_fields),
             'doc_form': document.doc_form