Fix return type

abetlen · Feb 11, 2024 · 48e53f5 · 48e53f5
1 parent 93ce456
commit 48e53f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ggml/experimental.py b/ggml/experimental.py
@@ -75,9 +75,9 @@ class Backend:
     def __init__(self, backend: ggml.ggml_backend_t):
         self.backend = backend
 
-    def alloc_ctx_tensors(self, ctx: Optional[Context] = None):
+    def alloc_ctx_tensors(self, ctx: Optional[Context] = None) -> "BackendBuffer":
         ctx = ctx or Context.current_context()
-        return ggml.ggml_backend_alloc_ctx_tensors(ctx.context, self.backend)
+        return BackendBuffer(ggml.ggml_backend_alloc_ctx_tensors(ctx.context, self.backend))
 
     @staticmethod
     def cpu():