Skip to content

Commit c53bbbc

Browse files
fix: fix missing register method params for framework models (#3122)
1 parent 8c93f92 commit c53bbbc

File tree

2 files changed

+11
-0
lines changed

2 files changed

+11
-0
lines changed

src/sagemaker/huggingface/model.py

+4
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,7 @@ def register(
304304
approval_status=None,
305305
description=None,
306306
drift_check_baselines=None,
307+
customer_metadata_properties=None,
307308
domain=None,
308309
):
309310
"""Creates a model package for creating SageMaker models or listing on Marketplace.
@@ -332,6 +333,8 @@ def register(
332333
or "PendingManualApproval". Defaults to ``PendingManualApproval``.
333334
description (str): Model Package description. Defaults to ``None``.
334335
drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
336+
customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
337+
metadata properties (default: None).
335338
domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
336339
"MACHINE_LEARNING" (default: None).
337340
@@ -362,6 +365,7 @@ def register(
362365
approval_status,
363366
description,
364367
drift_check_baselines=drift_check_baselines,
368+
customer_metadata_properties=customer_metadata_properties,
365369
domain=domain,
366370
)
367371

src/sagemaker/sklearn/model.py

+7
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,8 @@ def register(
151151
marketplace_cert=False,
152152
approval_status=None,
153153
description=None,
154+
drift_check_baselines=None,
155+
customer_metadata_properties=None,
154156
domain=None,
155157
):
156158
"""Creates a model package for creating SageMaker models or listing on Marketplace.
@@ -176,6 +178,9 @@ def register(
176178
approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
177179
or "PendingManualApproval" (default: "PendingManualApproval").
178180
description (str): Model Package description (default: None).
181+
drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
182+
customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
183+
metadata properties (default: None).
179184
domain (str): Domain values can be "COMPUTER_VISION", "NATURAL_LANGUAGE_PROCESSING",
180185
"MACHINE_LEARNING" (default: None).
181186
@@ -205,6 +210,8 @@ def register(
205210
marketplace_cert,
206211
approval_status,
207212
description,
213+
drift_check_baselines=drift_check_baselines,
214+
customer_metadata_properties=customer_metadata_properties,
208215
domain=domain,
209216
)
210217

0 commit comments

Comments
 (0)