Update app.py

app.py

@@ -1,35 +1,50 @@
 import gradio as gr
-import spaces  # 导入spaces时会应用torch补丁
+import spaces
 import torch
 
-spaces.zero.torch.unpatch()  # 手动调用解除torch补丁方法
 
-print('cuda avaliable: ' + str(torch.cuda.is_available()))  # false
+class LinearModel(torch.nn.Module):
+    def __init__(self):
+        super(LinearModel, self).__init__()
+        self.linear = torch.nn.Linear(1, 1, device='cuda')
 
-torch.set_default_device('cuda')
+    def forward(self, x):
+        out = self.linear(x)
+        print('weight device: ' + str(self.linear.weight.device))
+        return out
 
-zero = torch.Tensor([0]).cuda()  # 报错：无GPU资源
 
-one = torch.ones(2, 2, device='cuda')
-
-two = torch.matmul(one, one).cuda()
-
-
-def print_device():
+def process():
+    torch.set_default_device('cuda')
     print('cuda avaliable: ' + str(torch.cuda.is_available()))
-    print('zero device: ' + str(zero.device))
-    print('one device: ' + str(one.device))
-    print('two device: ' + str(two.device))
+    print('cuda device count: ' + str(torch.cuda.device_count()))
 
+    model = LinearModel().cuda()
+    x = torch.ones(1, device='cuda')
+    y = model(x)
 
-print_device()  # cpu
+    print(model)
+    print(x)
+    print(y)
 
 
 @spaces.GPU
 def greet(n):
-    print('on zero gpu')
-    print_device()  # cuda
-    return f"Hello {zero + n} Tensor"
+    process()
+    return f"Hello {n} Tensor"
+
+
+def func(n):
+    # step1, on cpu
+    process()
+
+    # step2, on gpu
+    res = greet(n)
+
+    # step3, on cpu
+    process()
+
+    return res
 
 
-gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text()).launch()
+gr.Interface(fn=func, inputs=gr.Number(), outputs=gr.Text()).launch()