Update app.py

app.py

@@ -1,54 +1,42 @@
 import gradio as gr
-import spaces
+import spaces  # 导入spaces时会应用torch补丁
 import torch
 import os
 
+spaces.zero.torch.unpatch()  # 手动调用解除torch补丁方法
+
+os.environ['CUDA_VISIBLE_DEVICES'] = 'MIG-2f70e35e-577e-52c1-9054-bc9f9d04054e'
+
+print('cuda visible devices: ' + str(os.getenv('CUDA_VISIBLE_DEVICES')))
+print('cuda avaliable: ' + str(torch.cuda.is_available()))
+print('cuda device count: ' + str(torch.cuda.device_count()))
+print('cuda device name: ' + str(torch.cuda.get_device_name()))
+print('cuda device capability: ' + str(torch.cuda.get_device_capability()))
+
 torch.set_default_device('cuda')
 
+zero = torch.Tensor([0]).cuda()
 
-class LinearModel(torch.nn.Module):
-    def __init__(self):
-        super(LinearModel, self).__init__()
-        self.linear = torch.nn.Linear(1, 1, device='cuda')
+one = torch.ones(2, 2, device='cuda')
 
-    def forward(self, x):
-        out = self.linear(x)
-        print('weight device: ' + str(self.linear.weight.device))
-        return out
+two = torch.matmul(one, one).cuda()
 
 
-def process(n: int):
-    print(f'\n===process step {n}')
-    print('cuda visible devices: ' + str(os.getenv('CUDA_VISIBLE_DEVICES')))
+def print_device():
     print('cuda avaliable: ' + str(torch.cuda.is_available()))
-    print('cuda device count: ' + str(torch.cuda.device_count()))
+    print('zero device: ' + str(zero.device))
+    print('one device: ' + str(one.device))
+    print('two device: ' + str(two.device))
 
-    model = LinearModel().cuda()
-    x = torch.ones(1, device='cuda')
-    y = model(x)
 
-    print(model)
-    print(x)
-    print(y)
+print_device()  # cpu
 
 
 @spaces.GPU
 def greet(n):
-    process(2)
-    return f"Hello {n} Tensor"
-
-
-def func(n):
-    # step1, on cpu
-    process(1)
-
-    # step2, on gpu
-    res = greet(n)
-
-    # step3, on cpu
-    process(3)
-
-    return res
+    print('on zero gpu')
+    print_device()  # cuda
+    return f"Hello {zero + n} Tensor"
 
 
-gr.Interface(fn=func, inputs=gr.Number(), outputs=gr.Text()).launch()
+gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text()).launch()