Add VRAM cleanup (gc + empty_cache) before training start
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -200,6 +200,17 @@ def _evaluate(
|
||||
return metrics
|
||||
|
||||
|
||||
def _clear_vram() -> None:
|
||||
"""Free VRAM from previous runs before starting."""
|
||||
import gc
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
allocated = torch.cuda.memory_allocated() / 1e9
|
||||
LOGGER.info("🧹 VRAM cleared. Current usage: %.2f GB", allocated)
|
||||
|
||||
|
||||
def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
"""Run full training loop."""
|
||||
coloredlogs.install(
|
||||
@@ -207,6 +218,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
|
||||
logger=LOGGER,
|
||||
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
|
||||
)
|
||||
_clear_vram()
|
||||
_set_seed(cfg.seed)
|
||||
output_dir = Path(cfg.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
Reference in New Issue
Block a user