fix: use callable map_location to skip Lightning model.to() — avoids torchmetrics CUDA assertion

app/models/tft_predictor.py

@@ -57,21 +57,6 @@ def _maybe_download(filename: str, local_path: str) -> bool:
         return False
 
 
-def _patch_checkpoint_for_cpu(obj):
-    """Recursively replace CUDA device references with CPU in a checkpoint dict."""
-    if isinstance(obj, dict):
-        return {k: _patch_checkpoint_for_cpu(v) for k, v in obj.items()}
-    if isinstance(obj, list):
-        return [_patch_checkpoint_for_cpu(v) for v in obj]
-    if isinstance(obj, torch.Tensor):
-        return obj.cpu()
-    if isinstance(obj, torch.device) and obj.type == "cuda":
-        return torch.device("cpu")
-    if isinstance(obj, str) and obj.lower().startswith("cuda"):
-        return "cpu"
-    return obj
-
-
 def load_model(model_path: str):
     """Load and cache the pytorch-forecasting TFT from a Lightning checkpoint."""
     global _model, _model_path_cached
@@ -84,20 +69,17 @@ def load_model(model_path: str):
 
     from pytorch_forecasting import TemporalFusionTransformer
 
-    # Checkpoint was saved on GPU (Colab). Patch the raw checkpoint to replace
-    # all CUDA device references before Lightning tries to move the model.
-    cpu_path = model_path + ".cpu"
-    if not os.path.exists(cpu_path):
-        ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
-        ckpt = _patch_checkpoint_for_cpu(ckpt)
-        torch.save(ckpt, cpu_path)
-
-    model = TemporalFusionTransformer.load_from_checkpoint(cpu_path, map_location="cpu")
-    model = model.cpu()
+    # Checkpoint was saved on GPU (Colab). Using a callable map_location moves
+    # all tensors to CPU AND skips Lightning's model.to(device) call, which is
+    # what triggers the torchmetrics CUDA assertion on CPU-only servers.
+    model = TemporalFusionTransformer.load_from_checkpoint(
+        model_path,
+        map_location=lambda storage, loc: storage.cpu(),
+    )
     model.eval()
     _model = model
     _model_path_cached = model_path
-    print(f"[tft] Loaded pytorch-forecasting TFT from {cpu_path}")
+    print(f"[tft] Loaded pytorch-forecasting TFT from {model_path}")
     return model