fix: Increase block size to reduce compile time BERT

pytorch · Aug 10, 2023 · b7d9d5a · b7d9d5a
1 parent 858d41d
commit b7d9d5a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 13 deletions.
diff --git a/tests/py/dynamo/converters/test_casts.py b/tests/py/dynamo/converters/test_casts.py
@@ -2,9 +2,7 @@
 import torch.nn as nn
 from harness import DispatchTestCase
 from torch.testing._internal.common_utils import run_tests
-from torch_tensorrt.dynamo.conversion.trt_interpreter import (
-    UnsupportedOperatorException,
-)
+from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
 
 
 class TestToCopyConverter(DispatchTestCase):

diff --git a/tests/py/dynamo/models/test_models.py b/tests/py/dynamo/models/test_models.py
@@ -1,18 +1,13 @@
-import torch
-import timm
-import pytest
 import unittest
 
+import pytest
+import timm
+import torch
 import torch_tensorrt as torchtrt
 import torchvision.models as models
-
+from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
 from transformers import BertModel
 
-from torch_tensorrt.dynamo.utils import (
-    COSINE_THRESHOLD,
-    cosine_similarity,
-)
-
 assertions = unittest.TestCase()
 
 
@@ -143,7 +138,7 @@ def test_bert_base_uncased(ir):
         "ir": ir,
         "pass_through_build_failures": True,
         "optimization_level": 1,
-        "min_block_size": 10,
+        "min_block_size": 15,
         "ir": "torch_compile",
     }
     trt_mod = torchtrt.compile(model, **compile_spec)