diff --git a/model/models/gemma2/model.go b/model/models/gemma2/model.go
index 2b8597c42..fbefebe2d 100644
--- a/model/models/gemma2/model.go
+++ b/model/models/gemma2/model.go
@@ -179,7 +179,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}
diff --git a/model/models/gemma3/model.go b/model/models/gemma3/model.go
index 32ad80f43..95f89ad48 100644
--- a/model/models/gemma3/model.go
+++ b/model/models/gemma3/model.go
@@ -150,7 +150,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}
diff --git a/model/models/llama/model.go b/model/models/llama/model.go
index 19a2ab8c4..87eb9b750 100644
--- a/model/models/llama/model.go
+++ b/model/models/llama/model.go
@@ -150,7 +150,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}
diff --git a/model/models/mllama/model.go b/model/models/mllama/model.go
index fa4d570ca..0aa11f178 100644
--- a/model/models/mllama/model.go
+++ b/model/models/mllama/model.go
@@ -154,7 +154,7 @@ func (m *Model) Forward(ctx ml.Context, opts input.Options) (ml.Tensor, error) {
 		return nil, err
 	}
 
-	outputs, err := ctx.Output().FromIntSlice(opts.Outputs, len(opts.Outputs))
+	outputs, err := ctx.Input().FromIntSlice(opts.Outputs, len(opts.Outputs))
 	if err != nil {
 		return nil, err
 	}