From 1cf5a567f0e20c58126728a1bc688a1f4c536976 Mon Sep 17 00:00:00 2001
From: Siyuan Feng <hzfengsy@sjtu.edu.cn>
Date: Fri, 30 Aug 2024 12:29:10 +0800
Subject: [PATCH] [Doc] Fix doc build error in e2e_opt_model.py

The `sys.exit` may stop the whole sphinx build process, but not the
single script execution.
---
 docs/how_to/tutorials/e2e_opt_model.py | 63 +++++++++++++-------------
 1 file changed, 31 insertions(+), 32 deletions(-)

diff --git a/docs/how_to/tutorials/e2e_opt_model.py b/docs/how_to/tutorials/e2e_opt_model.py
index a139e75cfe6a..0053d309d5a9 100644
--- a/docs/how_to/tutorials/e2e_opt_model.py
+++ b/docs/how_to/tutorials/e2e_opt_model.py
@@ -32,7 +32,6 @@
 # PyTorch.
 
 import os
-import sys
 import numpy as np
 import torch
 from torch import fx
@@ -101,39 +100,39 @@
 
 # Skip running in CI environment
 IS_IN_CI = os.getenv("CI", "") == "true"
-if IS_IN_CI:
-    sys.exit(0)
-
-with target:
-    mod = tvm.ir.transform.Sequential(
-        [
-            # Convert BatchNorm into a sequence of simpler ops for fusion
-            relax.transform.DecomposeOpsForInference(),
-            # Canonicalize the bindings
-            relax.transform.CanonicalizeBindings(),
-            # Run default optimization pipeline
-            relax.get_pipeline("zero"),
-            # Tune the model and store the log to database
-            relax.transform.MetaScheduleTuneIRMod({}, work_dir, TOTAL_TRIALS),
-            # Apply the database
-            relax.transform.MetaScheduleApplyDatabase(work_dir),
-        ]
-    )(mod)
-
-# Only show the main function
-mod["main"].show()
+if not IS_IN_CI:
+    with target:
+        mod = tvm.ir.transform.Sequential(
+            [
+                # Convert BatchNorm into a sequence of simpler ops for fusion
+                relax.transform.DecomposeOpsForInference(),
+                # Canonicalize the bindings
+                relax.transform.CanonicalizeBindings(),
+                # Run default optimization pipeline
+                relax.get_pipeline("zero"),
+                # Tune the model and store the log to database
+                relax.transform.MetaScheduleTuneIRMod({}, work_dir, TOTAL_TRIALS),
+                # Apply the database
+                relax.transform.MetaScheduleApplyDatabase(work_dir),
+            ]
+        )(mod)
+
+    # Only show the main function
+    mod["main"].show()
 
 ######################################################################
 # Build and Deploy
 # ----------------
 # Finally, we build the optimized model and deploy it to the target device.
-
-ex = relax.build(mod, target="cuda")
-dev = tvm.device("cuda", 0)
-vm = relax.VirtualMachine(ex, dev)
-# Need to allocate data and params on GPU device
-gpu_data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
-gpu_params = [tvm.nd.array(p, dev) for p in params["main"]]
-gpu_out = vm["main"](gpu_data, *gpu_params).numpy()
-
-print(gpu_out.shape)
+# We skip this step in the CI environment.
+
+if not IS_IN_CI:
+    ex = relax.build(mod, target="cuda")
+    dev = tvm.device("cuda", 0)
+    vm = relax.VirtualMachine(ex, dev)
+    # Need to allocate data and params on GPU device
+    gpu_data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype("float32"), dev)
+    gpu_params = [tvm.nd.array(p, dev) for p in params["main"]]
+    gpu_out = vm["main"](gpu_data, *gpu_params).numpy()
+
+    print(gpu_out.shape)