fixed the segmentation fault in CamemBERT and VIT

tancheng · Oct 9, 2024 · a72d3ac · a72d3ac
1 parent 06126cc
commit a72d3ac
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/experiments/CamemBERT/model/CamemBERT.py b/experiments/CamemBERT/model/CamemBERT.py
@@ -24,6 +24,7 @@ def __init__(self, model_name: str):
         self.layer3 = AutoModelForTokenClassification.from_pretrained(model_name).classifier
 
     def forward(self, input):
+        x = input.reshape(1, 7, 768)
         # Return only the logits.
         x = self.layer1(input).last_hidden_state
         x = self.layer2(x)
@@ -44,7 +45,10 @@ def forward(self, input):
 
 print("Parsing sentence tokens.")
 example_input = prepare_sentence_tokens(model_name, sentence)
-print(example_input.shape)
+print("example_input shape: ", example_input.shape)
+
+example_input = example_input.reshape(1, 7*768)
+print("example_input shape after reshaping: ", example_input.shape)
 
 print("Instantiating model.")
 model = OnlyLogitsHuggingFaceModel(model_name)

diff --git a/experiments/VIT/model/vit.py b/experiments/VIT/model/vit.py
@@ -25,6 +25,7 @@ def __init__(self):
 		self.layer2 = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').vit.layernorm
 		self.layer3 = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').classifier
 	def forward(self,x):
+		x = x.reshape(1, 197, 768)
 		x = self.layer1(x).last_hidden_state
 		x = self.layer2(x)
 		x = self.layer3(x)
@@ -37,6 +38,9 @@ def forward(self,x):
 model = prepare().eval()
 example_input = model(inputs)
 print(example_input.shape)
+example_input = example_input.reshape(1, 197*768)
+print("example_input shape after reshaping: ", example_input.shape)
+
 
 vit_model = vit().eval()
 output = vit_model(example_input)