Update customize.py - fix tensorrt _dev detection command (#913)

mlcommons · Aug 22, 2023 · 96297c3 · 96297c3
2 parents 64faa3d + 0cfe514
commit 96297c3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/cm-mlops/script/get-tensorrt/customize.py b/cm-mlops/script/get-tensorrt/customize.py
@@ -83,7 +83,10 @@ def preprocess(i):
         return {'return': 1, 'error': 'Windows is currently not supported!'}
 
     if env.get('CM_TENSORRT_TAR_FILE_PATH','')=='':
-        return {'return': 1, 'error': 'Please envoke cm run script "get tensorrt" --tar_file={full path to the TensorRT tar file}'}
+        tags = [ "get", "tensorrt" ]
+        if env.get('CM_TENSORRT_REQUIRE_DEV', '') != 'yes'
+            tags.append("_dev")
+        return {'return': 1, 'error': 'Please envoke cm run script ' + " ".join(tags) + " --tar_file={full path to the TensorRT tar file}'}
 
 
     print ('Untaring file - can take some time ...')