Merge pull request #8 from Yiran-ASU/yiran-fix3

Yiran fix3

experiments/CamemBERT/model/CamemBERT.py

@@ -46,13 +46,15 @@ def forward(self, input):
 print("Parsing sentence tokens.")
 example_input = prepare_sentence_tokens(model_name, sentence)
 print("example_input shape: ", example_input.shape)
+assert example_input.shape == (1, 7, 768), f"Expected shape (1,7,768), but got {example_input.shape}"
 
 # The original example_input shape is [1, 7, 768], now we reshape it into [1, 7*768]
 example_input = example_input.reshape(1, 7*768)
 print("example_input shape after reshaping: ", example_input.shape)
 
 print("Instantiating model.")
 model = OnlyLogitsHuggingFaceModel(model_name)
+assert example_input.shape == (1, 7*768), f"Expected the reshaped (1,7*768), but got {example_input.shape}"
 print(model(example_input).shape)
 
 linalg_on_tensors_mlir = torch_mlir.compile(

experiments/VIT/model/vit.py

@@ -38,13 +38,16 @@ def forward(self,x):
 model = prepare().eval()
 example_input = model(inputs)
 print(example_input.shape)
+assert example_input.shape == (1, 197, 768), f"Expected shape (1,197,768), but got {example_input.shape}"
+
 
 # The original example_input shape is [1, 197, 768], now we reshape it into [1, 197*768]
 example_input = example_input.reshape(1, 197*768)
 print("example_input shape after reshaping: ", example_input.shape)
 
 
 vit_model = vit().eval()
+assert example_input.shape == (1, 197*768), f"Expected the reshaped (1,197*768), but got {example_input.shape}"
 output = vit_model(example_input)
 print(output.shape)