Update app.py

app.py

@@ -16,7 +16,8 @@ class LinearModel(torch.nn.Module):
         return out
 
 
-def process():
+def process(n: int):
+    print(f'\n===process step {n}')
     print('cuda avaliable: ' + str(torch.cuda.is_available()))
     print('cuda device count: ' + str(torch.cuda.device_count()))
 
@@ -31,19 +32,19 @@ def process():
 
 @spaces.GPU
 def greet(n):
-    process()
+    process(2)
     return f"Hello {n} Tensor"
 
 
 def func(n):
     # step1, on cpu
-    process()
+    process(1)
 
     # step2, on gpu
     res = greet(n)
 
     # step3, on cpu
-    process()
+    process(3)
 
     return res