diff --git a/ml/src/trainers/tft.rs b/ml/src/trainers/tft.rs index 2b55e9ed3..8844bb3f4 100644 --- a/ml/src/trainers/tft.rs +++ b/ml/src/trainers/tft.rs @@ -1461,7 +1461,7 @@ impl TFTTrainer { None }; - for batch in val_loader.iter() { + for (i, batch) in val_loader.iter().enumerate() { // Convert batch to tensors let (static_tensor, hist_tensor, fut_tensor, target_tensor) = self.batch_to_tensors(batch)?; @@ -1493,6 +1493,16 @@ impl TFTTrainer { } batch_count += 1; + + // Clear CUDA cache every 10 batches to prevent accumulation (CRITICAL FIX) + if i % 10 == 0 && self.device.is_cuda() { + // Clear model's attention cache (prevents 2500MB leak during validation) + self.model.clear_cache(); + + if let Err(e) = Self::sync_cuda_device(&self.device) { + warn!("Failed to sync CUDA during validation batch {}: {}", i, e); + } + } } // Defensive check: if no validation batches, return zero loss (skip validation)