try

app.py

@@ -4,6 +4,21 @@ import os
 # 获取全部环境变量
 env_vars = os.environ
 
+os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = 0
+import torch
+
+if torch.cuda.is_available():
+    print("CUDA is available. Listing available GPUs:")
+    # 获取并打印GPU数量
+    num_gpus = torch.cuda.device_count()
+    for i in range(num_gpus):
+        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
+        # 其他相关信息，例如内存
+        print(f"  Memory Allocated: {torch.cuda.memory_allocated(i) / 1024 ** 2:.0f} MB")
+        print(f"  Memory Reserved: {torch.cuda.memory_reserved(i) / 1024 ** 2:.0f} MB")
+else:
+    print("CUDA is not available.")
+
 # 遍历并打印环境变量
 for key, value in env_vars.items():
     print(f"{key}: {value}")

spaces/zero/torch.py

@@ -241,14 +241,14 @@ if (torch := maybe_import_torch()):
         bitsandbytes.unpatch()
 
     def _move(nvidia_uuid: str):
-        os.environ['CUDA_VISIBLE_DEVICES'] = "5,8,10,11,14,15"
+        os.environ['CUDA_VISIBLE_DEVICES'] = "0,1,2,3,4,5"
         import subprocess
         subprocess.run(
             ['nvidia-smi'], text=True
         )
 
         print("_move(nvidia_uuid: str)")
-        torch.Tensor([5]).cuda() # CUDA init
+        torch.Tensor([0]).cuda() # CUDA init
         for op in to_ops.items():
             tensor, parsed_args = op
             _, dtype, _, memory_format = parsed_args