try

app.py

@@ -29,20 +29,28 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 ##
 # 检查CUDA是否可用
-os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-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.")
+def print_gpu():
+    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.")
 
 ##
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+print_gpu()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7,8'
+print_gpu()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7,8,9,10,11,12,13,14'
+print_gpu()
 
 
 if not torch.cuda.is_available():
@@ -67,19 +75,7 @@ def generate(
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
 
-    ## check inside function
-    # 检查CUDA是否可用
-    if torch.cuda.is_available():
-        print("check inside function: CUDA is available. Listing available GPUs:")
-        # 获取并打印GPU数量
-        num_gpus = torch.cuda.device_count()
-        for i in range(num_gpus):
-            print(f"check inside function: GPU {i}: {torch.cuda.get_device_name(i)}")
-            # 其他相关信息，例如内存
-            print(f"check inside function:   Memory Allocated: {torch.cuda.memory_allocated(i) / 1024 ** 2:.0f} MB")
-            print(f"check inside function:  Memory Reserved: {torch.cuda.memory_reserved(i) / 1024 ** 2:.0f} MB")
-    else:
-        print("check inside function: CUDA is not available.")
+    print_gpu()
 
     conversation = []
     if system_prompt: