setup_device.py

import torch

def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using GPU: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("CUDA not available. Using CPU.")
    return device

# Usage example:
if __name__ == "__main__":
    device = get_device()

    # Example: Create a tensor on the selected device
    x = torch.rand(3, 3).to(device)
    print(f"Tensor device: {x.device}")