From 142fe7edeb737acb375db04cc459738c7d62bde2 Mon Sep 17 00:00:00 2001
From: Alex Renda <alex@renda.org>
Date: Wed, 5 Jul 2023 19:34:19 -0700
Subject: [PATCH] tensor split

---
 llama_cpp/llama.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
index 688b2a74f0..ec2fec2864 100644
--- a/llama_cpp/llama.py
+++ b/llama_cpp/llama.py
@@ -207,6 +207,7 @@ def __init__(
         n_ctx: int = 512,
         n_parts: int = -1,
         n_gpu_layers: int = 0,
+        tensor_split: List[float] = [0],
         seed: int = 1337,
         f16_kv: bool = True,
         logits_all: bool = False,
@@ -254,6 +255,7 @@ def __init__(
         self.params = llama_cpp.llama_context_default_params()
         self.params.n_ctx = n_ctx
         self.params.n_gpu_layers = n_gpu_layers
+        self.params.tensor_split = tensor_split
         self.params.seed = seed
         self.params.f16_kv = f16_kv
         self.params.logits_all = logits_all