From 3fb723a1aa5c6eff216cb6d802f7af2e32088fc1 Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 2 Nov 2026 14:18:05 -0800
Subject: [PATCH 1/2] Keep cached class labels on their original device

---
 sdm/models/base/icl.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py
index 78734c8892..433bfad0ab 100644
--- a/sdm/models/base/icl.py
+++ b/sdm/models/base/icl.py
@@ -335,13 +335,19 @@ def fit(
                 )
 
             if x.is_cuda and len(contexts) > 1:
+                # Class labels are metadata; preserve their original device.
+                classes = batch_cache["classes"]
                 try:  # Copy to pinned CPU memory:
                     batch_cache = batch_cache._apply_tensor(
-                        lambda tensor: torch.ops.aten._to_copy.default(
-                            tensor,
-                            device="cpu",
-                            pin_memory=True,
-                            non_blocking=True,  # Required for `pin_memory`.
+                        lambda tensor, classes=classes: (
+                            tensor
+                            if tensor is classes
+                            else torch.ops.aten._to_copy.default(
+                                tensor,
+                                device="cpu",
+                                pin_memory=True,
+                                non_blocking=True,  # Needed for `pin_memory`.
+                            )
                         )
                     )
                 finally:

From 565fccb80cbffc15e12393c334ff2fed47bc457c Mon Sep 17 00:00:00 2001
From: Ard <88337265+aw471@users.noreply.github.com>
Date: Mon, 2 Nov 2026 14:20:55 -0800
Subject: [PATCH 2/2] Preserve class labels outside cache offloading

---
 sdm/models/base/icl.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py
index 433bfad0ab..2285346699 100644
--- a/sdm/models/base/icl.py
+++ b/sdm/models/base/icl.py
@@ -335,22 +335,19 @@ def fit(
                 )
 
             if x.is_cuda and len(contexts) > 1:
-                # Class labels are metadata; preserve their original device.
-                classes = batch_cache["classes"]
+                # Keep class labels on GPU for compiled prediction.
+                classes = batch_cache.pop("classes")
                 try:  # Copy to pinned CPU memory:
                     batch_cache = batch_cache._apply_tensor(
-                        lambda tensor, classes=classes: (
-                            tensor
-                            if tensor is classes
-                            else torch.ops.aten._to_copy.default(
-                                tensor,
-                                device="cpu",
-                                pin_memory=True,
-                                non_blocking=True,  # Needed for `pin_memory`.
-                            )
+                        lambda tensor: torch.ops.aten._to_copy.default(
+                            tensor,
+                            device="cpu",
+                            pin_memory=True,
+                            non_blocking=True,  # Required for `pin_memory`.
                         )
                     )
                 finally:
+                    batch_cache["classes"] = classes
                     torch.cuda.current_stream(x.device).synchronize()
 
             cache[i] = batch_cache