From 8b194b752042198c061f2f53797dfde5f9ac0d88 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 26 Feb 2025 12:16:59 -0800
Subject: [PATCH] kvcache: update tests

---
 kvcache/causal_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kvcache/causal_test.go b/kvcache/causal_test.go
index bbbdf8368..bd7d0ae8b 100644
--- a/kvcache/causal_test.go
+++ b/kvcache/causal_test.go
@@ -342,7 +342,7 @@ func (c *testContext) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
 	return out, nil
 }
 
-func (c *testContext) Forward(ml.Tensor) {}
+func (c *testContext) Forward(...ml.Tensor) ml.Context { return c }
 
 func (c *testContext) Compute(...ml.Tensor) {}