Skip to content

Commit 4643394

Browse files
fix: Fix role parameter for _RepackModelStep
1 parent f0efd27 commit 4643394

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

src/sagemaker/workflow/step_collections.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -336,7 +336,7 @@ def __init__(
336336
depends_on=depends_on,
337337
retry_policies=repack_model_step_retry_policies,
338338
sagemaker_session=estimator.sagemaker_session,
339-
role=estimator.sagemaker_session,
339+
role=estimator.role,
340340
model_data=model_data,
341341
entry_point=entry_point,
342342
source_dir=source_dir,
@@ -347,6 +347,7 @@ def __init__(
347347
description=description,
348348
display_name=display_name,
349349
repack_output_path=repack_output_path,
350+
**kwargs,
350351
)
351352
steps.append(repack_model_step)
352353
model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts

0 commit comments

Comments
 (0)