#include <cuda.h>
#include <iostream>
#include <cuda_runtime.h>

int main() {
    int cudaVersion;
    cudaError_t cudaResult = cudaRuntimeGetVersion(&cudaVersion);
    if (cudaResult != cudaSuccess) {
        printf("Failed to retrieve CUDA version: %s\n", cudaGetErrorString(cudaResult));
        return 1;
    }

    printf("CUDA Version: %d.%d\n", cudaVersion / 1000, (cudaVersion % 100) / 10);

    int driverVersion;
    cudaResult = cudaDriverGetVersion(&driverVersion);
    if (cudaResult != cudaSuccess) {
        printf("Failed to retrieve NVIDIA driver version: %s\n", cudaGetErrorString(cudaResult));
        return 1;
    }

    int major = driverVersion / 1000;
    int minor = (driverVersion % 100) / 10;
    printf("NVIDIA Driver Version: %d.%d\n", major, minor);

    int numDevices;
    cudaGetDeviceCount(&numDevices);

    if (numDevices == 0) {
        printf("No GPU device found. Make sure the container has GPU access.\n");
    } else {
        printf("%d GPU device(s) found. The container has GPU access.\n", numDevices);
    }

    while(1);

    return 0;
}