Skip to content

Commit ab9e309

Browse files
lanluo-nvidialanluo-nvidia
authored
fix ts fe test error. (#3619)
1 parent b0d5787 commit ab9e309

File tree

5 files changed

+77
-74
lines changed

5 files changed

+77
-74
lines changed

‎.github/scripts/install-torch-tensorrt.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ pip install --pre -r ${PWD}/tests/py/requirements.txt
1818
# eg. timm will install the latest torchvision, however we want to use the torchvision from nightly
1919
# reinstall torch torchvisionto make sure we have the correct version
2020
pip uninstall -y torch torchvision
21-
pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL} --no-deps
21+
pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL}
2222
pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}
2323

2424

‎MODULE.bazel‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ new_local_repository(
5252
new_local_repository(
5353
name = "cuda_win",
5454
build_file = "@//third_party/cuda:BUILD",
55-
path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8/",
55+
path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.9/",
5656
)
5757

5858
http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
@@ -66,7 +66,7 @@ http_archive(
6666
name = "libtorch",
6767
build_file = "@//third_party/libtorch:BUILD",
6868
strip_prefix = "libtorch",
69-
urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-shared-with-deps-latest.zip"],
69+
urls = ["https://download.pytorch.org/libtorch/nightly/cu129/libtorch-shared-with-deps-latest.zip"],
7070
)
7171

7272
# in aarch64 platform you can get libtorch via either local or wheel file
@@ -84,7 +84,7 @@ http_archive(
8484
name = "libtorch_win",
8585
build_file = "@//third_party/libtorch:BUILD",
8686
strip_prefix = "libtorch",
87-
urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-win-shared-with-deps-latest.zip"],
87+
urls = ["https://download.pytorch.org/libtorch/nightly/cu129/libtorch-win-shared-with-deps-latest.zip"],
8888
)
8989

9090
http_archive(

‎setup.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,6 @@ def build_libtorchtrt_cxx11_abi(
215215
print("Jetpack build")
216216

217217
if IS_SBSA:
218-
cmd.append("--platforms=//toolchains:sbsa")
219218
print("SBSA build")
220219

221220
if CI_BUILD:

‎tests/modules/hub.py‎

Lines changed: 67 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,14 @@
11
import importlib
2+
import importlib.util
23
import json
34
import os
45

56
import custom_models as cm
67
import torch
78

8-
if importlib.util.find_spec("torchvision"):
9-
import timm
10-
import torchvision.models as models
11-
129
torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
1310

11+
1412
torch_version = torch.__version__
1513

1614
# Detect case of no GPU before deserialization of models on GPU
@@ -22,7 +20,71 @@
2220
# Downloads all model files again if manifest file is not present
2321
MANIFEST_FILE = "model_manifest.json"
2422

25-
to_test_models = {}
23+
to_test_models = {
24+
"pooling": {"model": cm.Pool(), "path": "trace"},
25+
"module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
26+
"loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
27+
"loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
28+
"conditional": {"model": cm.FallbackIf(), "path": "script"},
29+
"inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
30+
"standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
31+
"tuple_input": {"model": cm.TupleInput(), "path": "script"},
32+
"list_input": {"model": cm.ListInput(), "path": "script"},
33+
"tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
34+
"list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
35+
"list_input_tuple_output": {
36+
"model": cm.ListInputTupleOutput(),
37+
"path": "script",
38+
},
39+
# "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
40+
}
41+
42+
if importlib.util.find_spec("torchvision"):
43+
import timm
44+
import torchvision.models as models
45+
46+
torchvision_models = {
47+
"alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
48+
"vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
49+
"squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
50+
"densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
51+
"inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
52+
"shufflenet": {
53+
"model": models.shufflenet_v2_x1_0(pretrained=True),
54+
"path": "both",
55+
},
56+
"mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
57+
"resnext50_32x4d": {
58+
"model": models.resnext50_32x4d(pretrained=True),
59+
"path": "both",
60+
},
61+
"wideresnet50_2": {
62+
"model": models.wide_resnet50_2(pretrained=True),
63+
"path": "both",
64+
},
65+
"mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
66+
"resnet18": {
67+
"model": torch.hub.load(
68+
"pytorch/vision:v0.9.0", "resnet18", pretrained=True
69+
),
70+
"path": "both",
71+
},
72+
"resnet50": {
73+
"model": torch.hub.load(
74+
"pytorch/vision:v0.9.0", "resnet50", pretrained=True
75+
),
76+
"path": "both",
77+
},
78+
"efficientnet_b0": {
79+
"model": timm.create_model("efficientnet_b0", pretrained=True),
80+
"path": "script",
81+
},
82+
"vit": {
83+
"model": timm.create_model("vit_base_patch16_224", pretrained=True),
84+
"path": "script",
85+
},
86+
}
87+
to_test_models.update(torchvision_models)
2688

2789

2890
def get(n, m, manifest):
@@ -77,67 +139,6 @@ def download_models(version_matches, manifest):
77139

78140

79141
def main():
80-
if not importlib.util.find_spec("torchvision"):
81-
print(f"torchvision is not installed, skip models download")
82-
return
83-
84-
to_test_models = {
85-
"alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
86-
"vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
87-
"squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
88-
"densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
89-
"inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
90-
"shufflenet": {
91-
"model": models.shufflenet_v2_x1_0(pretrained=True),
92-
"path": "both",
93-
},
94-
"mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
95-
"resnext50_32x4d": {
96-
"model": models.resnext50_32x4d(pretrained=True),
97-
"path": "both",
98-
},
99-
"wideresnet50_2": {
100-
"model": models.wide_resnet50_2(pretrained=True),
101-
"path": "both",
102-
},
103-
"mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
104-
"resnet18": {
105-
"model": torch.hub.load(
106-
"pytorch/vision:v0.9.0", "resnet18", pretrained=True
107-
),
108-
"path": "both",
109-
},
110-
"resnet50": {
111-
"model": torch.hub.load(
112-
"pytorch/vision:v0.9.0", "resnet50", pretrained=True
113-
),
114-
"path": "both",
115-
},
116-
"efficientnet_b0": {
117-
"model": timm.create_model("efficientnet_b0", pretrained=True),
118-
"path": "script",
119-
},
120-
"vit": {
121-
"model": timm.create_model("vit_base_patch16_224", pretrained=True),
122-
"path": "script",
123-
},
124-
"pooling": {"model": cm.Pool(), "path": "trace"},
125-
"module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
126-
"loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
127-
"loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
128-
"conditional": {"model": cm.FallbackIf(), "path": "script"},
129-
"inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
130-
"standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
131-
"tuple_input": {"model": cm.TupleInput(), "path": "script"},
132-
"list_input": {"model": cm.ListInput(), "path": "script"},
133-
"tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
134-
"list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
135-
"list_input_tuple_output": {
136-
"model": cm.ListInputTupleOutput(),
137-
"path": "script",
138-
},
139-
# "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
140-
}
141142

142143
manifest = None
143144
version_matches = False

‎tests/py/dynamo/automatic_plugin/test_flashinfer_rmsnorm.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
1-
import pytest
2-
3-
flashinfer = pytest.importorskip("flashinfer")
1+
import importlib
42
import unittest
53

4+
import pytest
65
import torch
76
import torch.nn as nn
87
import torch_tensorrt
@@ -12,6 +11,9 @@
1211

1312
from ..conversion.harness import DispatchTestCase
1413

14+
if importlib.util.find_spec("flashinfer"):
15+
import flashinfer
16+
1517

1618
@torch.library.custom_op("flashinfer::rmsnorm", mutates_args=()) # type: ignore[misc]
1719
def flashinfer_rmsnorm(
@@ -31,6 +33,7 @@ def _(input: torch.Tensor, weight: torch.Tensor, b: float = 1e-6) -> torch.Tenso
3133

3234

3335
@unittest.skip("Not Available")
36+
@unittest.skipIf(not importlib.util.find_spec("flashinfer"), "flashinfer not installed")
3437
class TestAutomaticPlugin(DispatchTestCase):
3538
@parameterized.expand(
3639
[

0 commit comments

Comments
 (0)