diff --git a/server/prompt_test.go b/server/prompt_test.go
index 6d04db53..21a5fd62 100644
--- a/server/prompt_test.go
+++ b/server/prompt_test.go
@@ -32,7 +32,7 @@ func TestChatPrompt(t *testing.T) {
 	mllamaModel := Model{Template: tmpl, ProjectorPaths: []string{"vision"}, Config: ConfigV2{ModelFamilies: []string{"mllama"}}}
 
 	createImg := func(width, height int) ([]byte, error) {
-		img := image.NewRGBA(image.Rect(0, 0, 5, 5))
+		img := image.NewRGBA(image.Rect(0, 0, width, height))
 		var buf bytes.Buffer
 
 		if err := png.Encode(&buf, img); err != nil {