Update app.py

app.py

@@ -1,5 +1,6 @@
+import spaces
 import torch
-device = f'cuda:0' if torch.cuda.is_available() else 'cpu'
+device = f'cuda' if torch.cuda.is_available() else 'cpu'
 print('device:', device)
 torch_version = torch.__version__
 print('torch_version:', torch_version)
@@ -7,8 +8,10 @@ print('torch_version:', torch_version)
 # 获取机器的CUDA版本
 # 在机器中输入命令nvidia-smi
 import subprocess
-try:
-    result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-    print('nvidia-smi output:', result.stdout)
-except Exception as e:
-    print(f'Error: {e}')
+@spaces.GPU
+def get_cuda_version():
+    try:
+        result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        print('nvidia-smi output:', result.stdout)
+    except Exception as e:
+        print(f'Error: {e}')