Update main.py

main.py

@@ -65,9 +65,9 @@ CONFIG_PATHS = {
     6: "config.yaml",  # Small v1 ONNX
     7: "config_version2.yaml",  # Small v2 ONNX
     8: "config_version3.yaml",  # Small v3 ONNX
-    9: "config.yaml"  # Medium v1 ONNX
+    9: "config.yaml",  # Medium v1 ONNX
     10: "config_version2.yaml",  # Medium v2 ONNX
-    11: "config_version3.yaml",  # Medium v3 ONNX
+    11: "config_version3.yaml"  # Medium v3 ONNX
 }
 
 # Mapping from PT index to ONNX index
@@ -76,8 +76,8 @@ PT_TO_ONNX_MAPPING = {
     1: 6,  # Small v2 -> ONNX
     2: 7,  # Medium v1 -> ONNX
     3: 8,  # Medium v2 -> ONNX
-    4: 9   # Medium v3 -> ONNX
-    5: 10   # Medium v3 -> ONNX
+    4: 9,   # Medium v3 -> ONNX
+    5: 10,   # Medium v3 -> ONNX
     6: 11   # Medium v3 -> ONNX
 }