Fix tests

Author: jackzhxng

examples/models/llava/test/test_llava.py

@@ -97,7 +97,7 @@ def test_llava_export(self):
         )[0]
         llava_module.run_method(
             "text_decoder",
-            (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_before_img),
+            (pte_embeds_before_img, torch.tensor([start_pos], dtype=torch.int64)),
         )
 
         # Update the start_pos. start_pos is used in kv cache. The source of truth
@@ -109,8 +109,8 @@ def test_llava_export(self):
         llava_module.run_method(
             "text_decoder",
             (
-                torch.tensor([start_pos], dtype=torch.int64),
                 pte_embeds_img,
+                torch.tensor([start_pos], dtype=torch.int64),
             ),
         )
 
@@ -123,7 +123,7 @@ def test_llava_export(self):
         )[0]
         pte_prefill_after_img = llava_module.run_method(
             "text_decoder",
-            (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_after_img),
+            (pte_embeds_after_img, torch.tensor([start_pos], dtype=torch.int64)),
         )[0]
 
         # Update the logits for each prefill (kv cache) step.
@@ -140,7 +140,7 @@ def test_llava_export(self):
             )[0]
             logits = llava_module.run_method(
                 "text_decoder",
-                (torch.tensor([start_pos + i], dtype=torch.int64), token_embeds),
+                (token_embeds, torch.tensor([start_pos + i], dtype=torch.int64)),
             )[0]
             new_tokens.append(torch.argmax(logits).item())