Update app.py

app.py

@@ -2,27 +2,27 @@ import gradio as gr
 import spaces
 import torch
 
+zero = torch.Tensor([0]).cuda()
 
-@spaces.GPU()
-def fn():
-    print('fn')
+one = torch.ones(2, 2, device='cuda')
 
+two = torch.matmul(one, one).cuda()
 
-def greet(n):
-    print('cuda avaliable: ' + str(torch.cuda.is_available()))
-
-    torch.set_default_device('cuda')
-
-    zero = torch.Tensor([0]).cuda()
-
-    one = torch.ones(2, 2, device='cuda')
-
-    two = torch.matmul(one, one).cuda()
 
+def print_device():
+    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_device()  # cpu
+
+
+@spaces.GPU
+def greet(n):
+    print('on zero gpu')
+    print_device()  # cuda
     return f"Hello {zero + n} Tensor"