Update ORT training dockerfile (#1733)

* update dockerfile & fix style * fix style with lower ruff and black --------- Co-authored-by: Jingya <[email protected]>
huggingface · Mar 1, 2024 · fd47a73 · fd47a73
1 parent ebb63f8
commit fd47a73
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 149 deletions.
diff --git a/examples/onnxruntime/training/docker/Dockerfile-ort1.14.1-cu116 b/examples/onnxruntime/training/docker/Dockerfile-ort1.14.1-cu116
diff --git a/examples/onnxruntime/training/docker/Dockerfile-ort1.15.1-cu118 b/examples/onnxruntime/training/docker/Dockerfile-ort1.15.1-cu118
diff --git a/...raining/docker/Dockerfile-ort1.16.3-cu118 → ...raining/docker/Dockerfile-ort1.17.1-cu118 b/...raining/docker/Dockerfile-ort1.16.3-cu118 → ...raining/docker/Dockerfile-ort1.17.1-cu118
@@ -65,7 +65,7 @@ RUN $PYTHON_EXE -m pip install onnx ninja
 RUN $PYTHON_EXE -m pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} -f https://download.pytorch.org/whl/${TORCH_CUDA_VERSION}
 
 # ORT Module
-RUN $PYTHON_EXE -m pip install onnxruntime-training==1.16.3 -f https://download.onnxruntime.ai/onnxruntime_stable_cu118.html
+RUN $PYTHON_EXE -m pip install onnxruntime-training==1.17.1 -f https://download.onnxruntime.ai/onnxruntime_stable_cu118.html
 RUN $PYTHON_EXE -m pip install torch-ort
 ENV TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX"
 RUN $PYTHON_EXE -m pip install --upgrade protobuf==3.20.2

diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py
@@ -918,9 +918,9 @@ def _from_pretrained(
             attribute_name_to_filename = {
                 "last_encoder_model_name": encoder_path.name,
                 "last_decoder_model_name": decoder_path.name if use_merged is False else None,
-                "last_decoder_with_past_model_name": decoder_with_past_path.name
-                if (use_merged is False and use_cache is True)
-                else None,
+                "last_decoder_with_past_model_name": (
+                    decoder_with_past_path.name if (use_merged is False and use_cache is True) else None
+                ),
                 "last_decoder_merged_name": decoder_merged_path.name if use_merged is True else None,
             }
             paths = {}

diff --git a/optimum/onnxruntime/runs/__init__.py b/optimum/onnxruntime/runs/__init__.py
@@ -65,9 +65,9 @@ def __init__(self, run_config):
             ref_keys=run_config["dataset"]["ref_keys"],
             task_args=run_config["task_args"],
             static_quantization=self.static_quantization,
-            num_calibration_samples=run_config["calibration"]["num_calibration_samples"]
-            if self.static_quantization
-            else None,
+            num_calibration_samples=(
+                run_config["calibration"]["num_calibration_samples"] if self.static_quantization else None
+            ),
             config=trfs_model.config,
             max_eval_samples=run_config["max_eval_samples"],
         )