core/device.py

import torch
import random
import numpy as np
from core.config import USE_GPU_IF_AVAILABLE, RANDOM_SEED

def set_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def get_device():
    if USE_GPU_IF_AVAILABLE and torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")

# Initialize system-wide settings
set_seed(RANDOM_SEED)
DEVICE = get_device()

print(f"[CORE] Using device: {DEVICE}")