Skip to content

Commit 135ffa7

Browse files
apivovarovlaurenyu
authored andcommitted
change: add Keras to the list of Neo-supported frameworks (#1177)
1 parent 042c3a0 commit 135ffa7

File tree

2 files changed

+5
-5
lines changed

2 files changed

+5
-5
lines changed

src/sagemaker/estimator.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -492,8 +492,8 @@ def compile_model(
492492
'var2':[1,1,28,28]}
493493
output_path (str): Specifies where to store the compiled model
494494
framework (str): The framework that is used to train the original
495-
model. Allowed values: 'mxnet', 'tensorflow', 'pytorch', 'onnx',
496-
'xgboost'
495+
model. Allowed values: 'mxnet', 'tensorflow', 'keras', 'pytorch',
496+
'onnx', 'xgboost'
497497
framework_version (str): The version of the framework
498498
compile_max_run (int): Timeout in seconds for compilation (default:
499499
3 * 60). After this amount of time Amazon SageMaker Neo

src/sagemaker/model.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
"qcs605",
4646
]
4747
)
48-
NEO_ALLOWED_FRAMEWORKS = set(["mxnet", "tensorflow", "pytorch", "onnx", "xgboost"])
48+
NEO_ALLOWED_FRAMEWORKS = set(["mxnet", "tensorflow", "keras", "pytorch", "onnx", "xgboost"])
4949

5050
NEO_IMAGE_ACCOUNT = {
5151
"us-west-1": "710691900526",
@@ -322,8 +322,8 @@ def compile(
322322
3 * 60). After this amount of time Amazon SageMaker Neo
323323
terminates the compilation job regardless of its current status.
324324
framework (str): The framework that is used to train the original
325-
model. Allowed values: 'mxnet', 'tensorflow', 'pytorch', 'onnx',
326-
'xgboost'
325+
model. Allowed values: 'mxnet', 'tensorflow', 'keras', 'pytorch',
326+
'onnx', 'xgboost'
327327
framework_version (str):
328328
329329
Returns:

0 commit comments

Comments
 (0)