cuda_check_in_ray.py

import ray
import torch
import os
import torch
print("CUDA Available:", torch.cuda.is_available())
print("CUDA Device Count:", torch.cuda.device_count())
print("Current CUDA Device:", torch.cuda.current_device())
print("CUDA Device Name:", torch.cuda.get_device_name(torch.cuda.current_device()))

print(os.environ.get("CUDA_VISIBLE_DEVICES"))

ray.init(num_gpus=1)

@ray.remote(num_gpus=1)
def check_cuda():
    print(os.environ.get("CUDA_VISIBLE_DEVICES"))
    return torch.cuda.is_available(), torch.cuda.device_count()

print(ray.get(check_cuda.remote()))