From ab2fd79759dfeff1dfc2f81af6fa44502a5d2070 Mon Sep 17 00:00:00 2001
From: Naomi Pentrel <5212232+npentrel@users.noreply.github.com>
Date: Thu, 18 Jul 2024 18:51:23 -0400
Subject: [PATCH 1/2] DOCS-2668: Update upload-training-script.md

---
 docs/services/ml/upload-training-script.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/docs/services/ml/upload-training-script.md b/docs/services/ml/upload-training-script.md
index 92f6d79bf6..4a548183a0 100644
--- a/docs/services/ml/upload-training-script.md
+++ b/docs/services/ml/upload-training-script.md
@@ -292,20 +292,20 @@ def save_tflite_classification(
     filename = os.path.join(model_dir, f"{model_name}.tflite")
     with open(filename, "wb") as f:
         f.write(tflite_model)
-```
+
 
 When you submit a training job with this training script, this function saves the model outputs to the `model_output_directory` in the cloud.
 Once the training job is complete, Viam looks at that directory and creates a package with all of the contents of the directory, creating or updating a registry item for the ML model.
 
 {{% /expand%}}
-{{%expand "Step 4: Write main function" %}}
+{{%expand "Step 5: Write main function" %}}
 
 Now, write all the code that runs the training job invoking the previously defined helper functions.
 Write this into the top level code of <file>training.py</file>, which is executed when the file runs as a script.
 
 For the [example classification training script](https://app.viam.com/packages/e76d1b3b-0468-4efd-bb7f-fb1d2b352fcb/custom-training-classification/ml_training/latest/e76d1b3b-0468-4efd-bb7f-fb1d2b352fcb) that trains a classification model using TensorFlow and Keras, `__main__` looks like this:
 
-```python {class="line-numbers linkable-line-numbers"}
+python {class="line-numbers linkable-line-numbers"}
 if __name__ == "__main__":
     DATA_JSON, MODEL_DIR = parse_args()
     # Set up compute device strategy. If GPUs are available, they will be used
@@ -389,7 +389,7 @@ You can use the [`viam training-script upload`](/cli/#training-script) command t
 
 For example:
 
-```sh {class="line-numbers linkable-line-numbers"}
+```sh {class="command-line" data-prompt="$"}
 viam training-script upload --path=<path-to-tar.gz> --org-id=<INSERT ORG ID> --script-name="MyCustomTrainingScript"
 ```
 
@@ -411,7 +411,7 @@ For example:
 {{< tabs >}}
 {{% tab name="from-registry" %}}
 
-```sh {class="line-numbers linkable-line-numbers"}
+```sh {class="command-line" data-prompt="$"}
 viam train submit custom from-registry --dataset-id=<INSERT DATASET ID> --org-id=<INSERT ORG ID> --model-name="MyRegistryModel" --model-version="2" --version="1" --path=<path-to-tar.gz> --script-name="MyCustomTrainingScript"
 ```
 
@@ -420,7 +420,7 @@ This command submits a training job to the previously uploaded `"MyCustomTrainin
 {{% /tab %}}
 {{% tab name="with-upload" %}}
 
-```sh {class="line-numbers linkable-line-numbers"}
+```sh {class="command-line" data-prompt="$"}
 viam train submit custom with-upload --dataset-id=<INSERT DATASET ID> --model-org-id=<INSERT ORG ID> --model-name="MyRegistryModel" --model-type="single_label_classification" --model-version="2" --version="1" --path=<path-to-tar.gz> --script-name="MyCustomTrainingScript"
 ```
 

From 556b067f2bb6147a1c39f86f5b67409434206ee7 Mon Sep 17 00:00:00 2001
From: Naomi Pentrel <5212232+npentrel@users.noreply.github.com>
Date: Thu, 18 Jul 2024 18:52:20 -0400
Subject: [PATCH 2/2] Update upload-training-script.md

---
 docs/services/ml/upload-training-script.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/services/ml/upload-training-script.md b/docs/services/ml/upload-training-script.md
index 4a548183a0..8012c1c783 100644
--- a/docs/services/ml/upload-training-script.md
+++ b/docs/services/ml/upload-training-script.md
@@ -292,7 +292,7 @@ def save_tflite_classification(
     filename = os.path.join(model_dir, f"{model_name}.tflite")
     with open(filename, "wb") as f:
         f.write(tflite_model)
-
+```
 
 When you submit a training job with this training script, this function saves the model outputs to the `model_output_directory` in the cloud.
 Once the training job is complete, Viam looks at that directory and creates a package with all of the contents of the directory, creating or updating a registry item for the ML model.
@@ -305,7 +305,7 @@ Write this into the top level code of <file>training.py</file>, which is execute
 
 For the [example classification training script](https://app.viam.com/packages/e76d1b3b-0468-4efd-bb7f-fb1d2b352fcb/custom-training-classification/ml_training/latest/e76d1b3b-0468-4efd-bb7f-fb1d2b352fcb) that trains a classification model using TensorFlow and Keras, `__main__` looks like this:
 
-python {class="line-numbers linkable-line-numbers"}
+```python {class="line-numbers linkable-line-numbers"}
 if __name__ == "__main__":
     DATA_JSON, MODEL_DIR = parse_args()
     # Set up compute device strategy. If GPUs are available, they will be used