[Bugfix] Use platform-agnostic device in example_connector load (#42926)

Signed-off-by: Revital Sur <eres@il.ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-05-19 03:12:04 +00:00
parent 287471b994
commit afd7b1dce9
1 changed files with 3 additions and 1 deletions
@@ -188,7 +188,9 @@ class ExampleConnector(KVConnectorBase_V1):
                filename = self._generate_filename_debug(
                    layer_name, request.token_ids, request.mm_hashes
                )
-                kv_cache = safetensors.torch.load_file(filename)["kv_cache"].cuda()
+                kv_cache = safetensors.torch.load_file(
+                    filename, device=str(kv_cache_layer.device)
+                )["kv_cache"]
                if isinstance(attn_metadata, dict):
                    inject_kv_into_layer(
                        kv_cache_layer,