"""
GPU-optimized training script - ensures GPU is used and optimized.
"""

import torch
import os
import sys

# Force CUDA
torch.cuda.empty_cache()
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from training.train_with_eval import main
import argparse

if __name__ == '__main__':
    # Verify GPU
    if not torch.cuda.is_available():
        print("ERROR: GPU not available! Exiting.")
        print("This script requires GPU. Check CUDA installation.")
        sys.exit(1)
    
    print("="*70)
    print("GPU TRAINING SESSION")
    print("="*70)
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"CUDA Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
    print(f"CUDA Version: {torch.version.cuda}")
    print("="*70)
    print()
    
    # Call main training function
    main()