Skip to content

Commit 565af7c

Browse files
committed
update
1 parent b9d0e27 commit 565af7c

File tree

1 file changed

+2
-4
lines changed

1 file changed

+2
-4
lines changed

src/sagemaker/model.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -620,15 +620,13 @@ def _compilation_job_config(
620620
}
621621

622622
if (
623-
framework.lower() == "pytorch"
623+
framework.lower() == "pytorch" or framework.lower() == "tensorflow"
624+
and target_instance_type is not None
624625
and re.match("(?=^ml_)(?!ml_inf)", target_instance_type) is not None
625626
and framework_version is not None
626627
):
627628
input_model_config["FrameworkVersion"] = utils.get_short_version(framework_version)
628629

629-
if framework.lower() == "tensorflow" and framework_version is not None:
630-
input_model_config["FrameworkVersion"] = framework_version
631-
632630
role = self.sagemaker_session.expand_role(role)
633631
output_model_config = {
634632
"S3OutputLocation": output_path,

0 commit comments

Comments
 (0)