TorchScript and Deployment

## Why a checkpoint is device-flavoured When you `torch.save` a tensor, PyTorch stores its data and the device it lived on. Restoring it means re-materializing the storage on that device. On a CPU-only inference box there is no `cuda:0` to restore to, so `torch.load` raises rather than guessing — the message tells you to use `map_location`, and it names `torch.cuda.is_available() is False` as the reason. ## map_location `map_location` tells the deserializer where storages should land. The common forms: - `map_location="cpu"` — everything to CPU. Safest default for a load-then-move flow. - `map_location=torch.device("cuda:0")` — remap onto a specific device, useful when the checkpoint came from `cuda:3` and this box has one GPU. - `map_location=lambda storage, loc: storage` — the old idiom meaning "leave it where the current process would put it", i.e. CPU. The standard deployment pattern is: load to CPU, `load_state_dict`, then `model.to(device)` once at startup. That works identically whether the target has a GPU or not, and it is the version you should be able to write on a whiteboard. ## Devices must agree at run time Parameters and inputs must be on the same device or the forward pass raises "Expected all tensors to be on the same device". So a serving wrapper computes `device` once, moves the model there once, and moves each incoming batch there per request. Do not scatter `.cuda()` calls through model code — a hardcoded `.cuda()` makes the module unloadable on a CPU box no matter what `map_location` you passed. ## state_dict versus the whole module `torch.save(model, path)` pickles the module object, which stores a reference to your class by import path. Load it on a machine where that module has moved, been renamed or has a different signature and it breaks. `torch.save(model.state_dict(), path)` stores only a dict of tensor names to tensors; you reconstruct the architecture in code and call `load_state_dict`. That is the format to use for anything that outlives the training script. `load_state_dict` is strict by default: unexpected or missing keys raise. Common causes are a `module.` prefix from `DataParallel`/DDP-wrapped training, and an `_orig_mod.` prefix from a `torch.compile` wrapper. The fix is to strip the prefix when saving or loading rather than to pass `strict=False`, which silently leaves layers at their random initialization — a genuinely dangerous escape hatch. ## The weights_only default Since torch 2.6, `torch.load` defaults to `weights_only=True`. Unpickling is code execution: a maliciously crafted checkpoint can run arbitrary Python the moment you load it, which matters the instant you accept model files from a hub, a customer or a shared bucket. `weights_only=True` restricts deserialization to tensors and a small allowlist of safe types. The practical consequence is that old checkpoints containing non-tensor objects — a whole pickled module, a custom optimizer config class, a numpy scalar wrapper — now fail to load. The right responses, in order: re-save as a plain `state_dict`; or register the specific classes you trust with `torch.serialization.add_safe_globals`. Setting `weights_only=False` is only defensible for a file you produced and control, and saying that out loud is what separates a good answer from a shrug. ## What to check once it loads Put the model in evaluation mode before serving and confirm parameter dtypes are what you expect — a checkpoint saved from mixed-precision training may hold half-precision tensors that a CPU box will run slowly or refuse for some operators. And run one known input through both the training-time model and the freshly loaded one, comparing outputs with a tolerance. A checkpoint that loads without error is not the same as a checkpoint that loaded correctly, and `strict=False` plus a prefix mismatch produces exactly that failure: a service that starts fine and returns garbage.