Skip to content

Commit 38b9114

Browse files
committed
fix: fix missing register method params for framework models
1 parent 255a339 commit 38b9114

File tree

2 files changed

+11
-0
lines changed

2 files changed

+11
-0
lines changed

src/sagemaker/huggingface/model.py

Lines changed: 4 additions & 0 deletions
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
):
308309
"""Creates a model package for creating SageMaker models or listing on Marketplace.
309310
@@ -331,6 +332,8 @@ def register(
331332
or "PendingManualApproval". Defaults to ``PendingManualApproval``.
332333
description (str): Model Package description. Defaults to ``None``.
333334
drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
335+
customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
336+
metadata properties (default: None).
334337
335338
Returns:
336339
A `sagemaker.model.ModelPackage` instance.
@@ -359,6 +362,7 @@ def register(
359362
approval_status,
360363
description,
361364
drift_check_baselines=drift_check_baselines,
365+
customer_metadata_properties=customer_metadata_properties,
362366
)
363367

364368
def prepare_container_def(

src/sagemaker/sklearn/model.py

Lines changed: 7 additions & 0 deletions
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
):
155157
"""Creates a model package for creating SageMaker models or listing on Marketplace.
156158
@@ -175,6 +177,9 @@ def register(
175177
approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
176178
or "PendingManualApproval" (default: "PendingManualApproval").
177179
description (str): Model Package description (default: None).
180+
drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
181+
customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
182+
metadata properties (default: None).
178183
179184
Returns:
180185
A `sagemaker.model.ModelPackage` instance.
@@ -202,6 +207,8 @@ def register(
202207
marketplace_cert,
203208
approval_status,
204209
description,
210+
drift_check_baselines=drift_check_baselines,
211+
customer_metadata_properties=customer_metadata_properties,
205212
)
206213

207214
def prepare_container_def(

0 commit comments

Comments
 (0)