fix: Torch nightly version 2

Author: gs-olive

tests/modules/custom_models.py

@@ -1,8 +1,9 @@
+from typing import Dict, List, Tuple
+
 import torch
 import torch.nn as nn
-from transformers import BertModel, BertTokenizer, BertConfig
 import torch.nn.functional as F
-from typing import Tuple, List, Dict
+from transformers import BertConfig, BertModel, BertTokenizer
 
 
 # Sample Pool Model (for testing plugin serialization)
@@ -174,16 +175,19 @@ def BertModule():
     segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
     tokens_tensor = torch.tensor([indexed_tokens])
     segments_tensors = torch.tensor([segments_ids])
-    config = BertConfig(
-        vocab_size_or_config_json_file=32000,
-        hidden_size=768,
-        num_hidden_layers=12,
-        num_attention_heads=12,
-        intermediate_size=3072,
-        torchscript=True,
-    )
-    model = BertModel(config)
+
+    model_kwargs = {
+        "vocab_size_or_config_json_file": 32000,
+        "hidden_size": 768,
+        "num_hidden_layers": 12,
+        "num_attention_heads": 12,
+        "intermediate_size": 3072,
+        "use_cache": False,
+        "output_attentions": False,
+        "output_hidden_states": False,
+        "torchscript": True,
+    }
+    model = BertModel.from_pretrained(model_name, **model_kwargs)
     model.eval()
-    model = BertModel.from_pretrained(model_name, torchscript=True)
     traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
     return traced_model