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}")