leihanchen

import argparse
import os

import torch

# ---Configuration---
# Default path can still be overridden via --file at runtime.
FILE_PATH = "model.pth"

leihanchen

import torch

# Check if CUDA is available
if torch.cuda.is_available():
    # Get the version of CUDA that PyTorch is using
    cuda_version = torch.version.cuda
    print(f"PyTorch CUDA version: {cuda_version}")

    # Get the number of available CUDA devices
    num_devices = torch.cuda.device_count()