diff --git a/magi_compiler/_api.py b/magi_compiler/_api.py
index e7c6209..7bf5f21 100644
--- a/magi_compiler/_api.py
+++ b/magi_compiler/_api.py
@@ -509,6 +509,17 @@ def _cpu_apply(self, fn):
         id_cpu_lambda = getattr(fn, "__qualname__", "") == "Module.cpu.<locals>.<lambda>"
         is_to_lambda = getattr(fn, "__qualname__", "") == "Module.to.<locals>.convert"
 
+        # Detect .to("cuda:X") by probing the lambda on a small CPU tensor.
+        # .to("cpu") also produces is_to_lambda=True but should not trigger offload.
+        is_to_cuda = False
+        if is_to_lambda and not getattr(self, "_magi_offloaded_once", False):
+            try:
+                is_to_cuda = fn(torch.empty(0, device="cpu")).is_cuda
+            except Exception:
+                pass
+
+        is_moving_to_gpu = is_cuda_lambda or is_to_cuda
+
         # after first time to call _apply(cuda), skip "Module.to" and "Module.cpu" and "Module.cuda"
         if getattr(self, "_magi_offloaded_once", False):
             if is_cuda_lambda or id_cpu_lambda or is_to_lambda:
@@ -516,8 +527,8 @@ def _cpu_apply(self, fn):
             else:
                 return _orig_apply(self, fn)
         else:
-            # first time to call _apply(cuda), move all parameters/buffers to CPU
-            if not is_cuda_lambda:
+            # first time to call _apply(cuda) or _apply(to_cuda), move all parameters/buffers to CPU
+            if not is_moving_to_gpu:
                 return _orig_apply(self, fn)
 
         # move all parameters/buffers to CPU
diff --git a/magi_compiler/config.py b/magi_compiler/config.py
index f0206ac..d5592d8 100644
--- a/magi_compiler/config.py
+++ b/magi_compiler/config.py
@@ -180,6 +180,7 @@ class OffloadConfig(BaseModel):
         OffloadPolicy.COST_EFFECTIVE, description="The policy for offloading the model to CPU."
     )
     bandwidth_safety_factor: float = Field(0.9, description="The safety factor for the H2D bandwidth.")
+    max_prefetch_lookahead: int = Field(2, description="Max layers to prefetch ahead. 0 disables prefetch to save GPU memory.")
 
 
 class FSDPConfig(BaseModel):
diff --git a/magi_compiler/offload/scheduler.py b/magi_compiler/offload/scheduler.py
index 566a4ea..c3569b4 100644
--- a/magi_compiler/offload/scheduler.py
+++ b/magi_compiler/offload/scheduler.py
@@ -154,7 +154,7 @@ def prefetch(self, current_node_name: str, ctx: OffloadRuntimeContext):
         except ValueError:
             return
 
-        max_lookahead = 2
+        max_lookahead = self.compile_config.offload_config.max_prefetch_lookahead
         target_node = None
         is_next_iter = False