<GENERATED-CODE>
1. Update `ml/src/tft/mod.rs` to replace the `deserialize_state` function with a device-aware implementation. This new version uses `candle_core::safetensors::load(&path, &self.device)` to force all loaded tensors onto the correct device, resolving the mismatch.

<UPDATED_EXISTING_FILE: ml/src/tft/mod.rs>
```rust
// ... (code before line 975)

    async fn deserialize_state(&mut self, data: &[u8]) -> Result<(), MLError> {
        // Write bytes to a temporary file to leverage `safetensors::load`.
        let temp_dir = std::env::temp_dir();
        let temp_path = temp_dir.join(format!("tft_restore_{}.safetensors", Uuid::new_v4()));

        std::fs::write(&temp_path, data)
            .map_err(|e| MLError::ModelError(format!("Failed to write temp checkpoint: {}", e)))?;

        // CORE FIX: Load all tensors from the file, forcing them onto the model's
        // configured device (`self.device`). This prevents device mismatches when
        // loading a checkpoint from a different environment (e.g., GPU -> CPU).
        let tensors = candle_core::safetensors::load(&temp_path, &self.device)
            .map_err(|e| MLError::ModelError(format!("Failed to load tensors with device override: {}", e)))?;

        // Clean up the temporary file immediately.
        let _ = std::fs::remove_file(&temp_path);

        // Get mutable access to the VarMap to update the variables.
        let varmap_mut = Arc::get_mut(&mut self.varmap).ok_or_else(|| {
            MLError::ModelError(
                "Cannot load checkpoint: VarMap has multiple references. \
                 This indicates the model is being shared across threads. \
                 Clone the model before loading checkpoint."
                    .to_string(),
            )
        })?;

        // Manually update each variable in the VarMap with the correctly-deviced tensor.
        // This logic replaces the opaque `varmap.load()` with an explicit, device-aware update loop.
        for (name, tensor) in tensors.into_iter() {
            if let Some(var) = varmap_mut.get_mut(&name) {
                var.set(&tensor).map_err(|e| {
                    MLError::ModelError(format!("Failed to set tensor for var '{}': {}", name, e))
                })?;
            } else {
                warn!(
                    "Tensor '{}' found in checkpoint but not in model's VarMap. This can happen if the model architecture has changed.",
                    name
                );
            }
        }

        debug!(
            "Deserialized TFT state from {} bytes to device {:?}",
            data.len(),
            self.device
        );
        Ok(())
    }

// ... (code after line 1008)
```
</UPDATED_EXISTING_FILE>
</GENERATED-CODE>