[CI] Update GPU image for PyTorch 1.11 (#10849)

* Update Jenkinsfile to point to the new GPU image * fix onnx test * fixed keras tutorial
uwsampl · Apr 1, 2022 · 621f777 · 621f777
1 parent 95df0eb
commit 621f777
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -47,7 +47,7 @@
 import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
 // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. -->
 ci_lint = 'tlcpack/ci-lint:v0.69'
-ci_gpu = 'tlcpack/ci-gpu:v0.82'
+ci_gpu = 'tlcpack/ci-gpu:v0.83'
 ci_cpu = 'tlcpack/ci-cpu:v0.82'
 ci_wasm = 'tlcpack/ci-wasm:v0.72'
 ci_i386 = 'tlcpack/ci-i386:v0.75'

diff --git a/gallery/how_to/compile_models/from_keras.py b/gallery/how_to/compile_models/from_keras.py
@@ -39,6 +39,7 @@
 import tvm.relay as relay
 from tvm.contrib.download import download_testdata
 import keras
+import tensorflow as tf
 import numpy as np
 
 ######################################################################
@@ -65,7 +66,7 @@
 
 
 weights_path = download_testdata(weights_url, weights_file, module="keras")
-keras_resnet50 = keras.applications.resnet50.ResNet50(
+keras_resnet50 = tf.keras.applications.resnet50.ResNet50(
     include_top=True, weights=None, input_shape=(224, 224, 3), classes=1000
 )
 keras_resnet50.load_weights(weights_path)
@@ -76,7 +77,7 @@
 # A single cat dominates the examples!
 from PIL import Image
 from matplotlib import pyplot as plt
-from keras.applications.resnet50 import preprocess_input
+from tensorflow.keras.applications.resnet50 import preprocess_input
 
 img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
 img_path = download_testdata(img_url, "cat.png", module="data")

diff --git a/jenkins/Jenkinsfile.j2 b/jenkins/Jenkinsfile.j2
@@ -49,7 +49,7 @@ import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
 
 // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. -->
 ci_lint = 'tlcpack/ci-lint:v0.69'
-ci_gpu = 'tlcpack/ci-gpu:v0.82'
+ci_gpu = 'tlcpack/ci-gpu:v0.83'
 ci_cpu = 'tlcpack/ci-cpu:v0.82'
 ci_wasm = 'tlcpack/ci-wasm:v0.72'
 ci_i386 = 'tlcpack/ci-i386:v0.75'

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
@@ -2993,7 +2993,7 @@ def forward(self, input):
         torch.onnx.export(layer, dummy_input, file_name, export_params=True)
 
         onnx_model = onnx.load(file_name)
-        relay.frontend.from_onnx(onnx_model, {"0": input_size})
+        relay.frontend.from_onnx(onnx_model, {"onnx::Reshape_0": input_size})
 
 
 @tvm.testing.parametrize_targets