Skip to content

Commit 26d1051

Browse files
committed
fix: Increase block size to reduce compile time BERT
1 parent 31c1cfd commit 26d1051

File tree

2 files changed

+6
-13
lines changed

2 files changed

+6
-13
lines changed

tests/py/dynamo/converters/test_casts.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@
22
import torch.nn as nn
33
from harness import DispatchTestCase
44
from torch.testing._internal.common_utils import run_tests
5-
from torch_tensorrt.dynamo.conversion.trt_interpreter import (
6-
UnsupportedOperatorException,
7-
)
5+
from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
86

97

108
class TestToCopyConverter(DispatchTestCase):

tests/py/dynamo/models/test_models.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,13 @@
1-
import torch
2-
import timm
3-
import pytest
41
import unittest
52

3+
import pytest
4+
import timm
5+
import torch
66
import torch_tensorrt as torchtrt
77
import torchvision.models as models
8-
8+
from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
99
from transformers import BertModel
1010

11-
from torch_tensorrt.dynamo.utils import (
12-
COSINE_THRESHOLD,
13-
cosine_similarity,
14-
)
15-
1611
assertions = unittest.TestCase()
1712

1813

@@ -143,7 +138,7 @@ def test_bert_base_uncased(ir):
143138
"ir": ir,
144139
"pass_through_build_failures": True,
145140
"optimization_level": 1,
146-
"min_block_size": 10,
141+
"min_block_size": 15,
147142
"ir": "torch_compile",
148143
}
149144
trt_mod = torchtrt.compile(model, **compile_spec)

0 commit comments

Comments
 (0)