[llava] Remove torch.jit.save in llava example

Summary: As titled.

Test Plan: Rely on CI test

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

.ci/scripts/test_llava.sh

@@ -149,13 +149,8 @@ run_and_verify() {
 
     # verify result.txt
     RESULT=$(cat result.txt)
-    # set the expected prefix to be the same as prompt because there's a bug in sdpa_with_kv_cache that causes <unk> tokens.
-    if [[ "$(uname)" == "Darwin" ]]; then
-        EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with several players on the court. One of the players is dribbling the ball, while the others are in various"
-    else
-        # set the expected prefix to be the same as prompt because there's a bug in sdpa_with_kv_cache that causes <unk> tokens.
-        EXPECTED_PREFIX="ASSISTANT: image"
-    fi
+    EXPECTED_PREFIX="ASSISTANT: image captures a basketball game in progress, with several players on the court. One of the players is dribbling the ball, while"
+
     if [[ "${RESULT}" == *"${EXPECTED_PREFIX}"* ]]; then
         echo "Expected result prefix: ${EXPECTED_PREFIX}"
         echo "Actual result: ${RESULT}"

examples/models/llava/image_util.py

@@ -48,12 +48,7 @@ def prepare_image(image: Image, target_h: int, target_w: int) -> torch.Tensor:
 
 def serialize_image(image: torch.Tensor, path: str) -> None:
     copy = torch.tensor(image)
-    m = nn.Module()
-    par = nn.Parameter(copy, requires_grad=False)
-    m.register_parameter("0", par)
-    tensors = torch.jit.script(m)
-    tensors.save(path)
-
+    torch.save(copy, path)
     logging.info(f"Saved image tensor to {path}")