diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py
index 78734c8892..2285346699 100644
--- a/sdm/models/base/icl.py
+++ b/sdm/models/base/icl.py
@@ -335,6 +335,8 @@ def fit(
                 )
 
             if x.is_cuda and len(contexts) > 1:
+                # 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: torch.ops.aten._to_copy.default(
@@ -345,6 +347,7 @@ def fit(
                         )
                     )
                 finally:
+                    batch_cache["classes"] = classes
                     torch.cuda.current_stream(x.device).synchronize()
 
             cache[i] = batch_cache