#!/bin/bash

# Set CUDA environment variables
export CUDA_HOME=/usr/local/cuda
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

# Create symlink if it doesn't exist
ln -sf /usr/local/cuda-12.4 /usr/local/cuda

# Install NVIDIA toolkit and dependencies
apt-get update && apt-get install -y \
    cuda-toolkit \
    ninja-build \
    build-essential

# Install Python dependencies in the correct order
pip install torch==2.4.0 torchvision torchaudio
pip install -r requirements.txt