diff --git a/src/llama.cpp b/src/llama.cpp
index 05755936..f462b225 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -1640,7 +1640,7 @@ static void restore_recurrent_cache_tensors(int step, ggml_backend_sched_t sched
         size_t ssm_bytes, size_t conv_bytes,
         ggml_tensor * s_l, ggml_tensor * per_step_ssm, ggml_tensor * per_step_conv,
         std::unordered_set<ggml_backend_t> & backends_to_sync) {
-    auto dst_backend = ggml_backend_sched_get_backend(sched, ggml_backend_sched_get_backend_idx(sched, s_l->buffer));
+    auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, s_l);
     auto dst = *s_l;
     dst.ne[0] = ssm_bytes/sizeof(float);
     dst.nb[1] = dst.nb[2] = dst.nb[3] = ssm_bytes + conv_bytes;