Update app.py

app.py

@@ -6,20 +6,25 @@ print('cuda avaliable: ' + str(torch.cuda.is_available()))
 torch.set_default_device('cuda')
 
 zero = torch.Tensor([0]).cuda()
-print('zero device: ' + str(zero.device))  # cpu
 
 one = torch.ones(2, 2, device='cuda')
-print('one device: ' + str(one.device))  # cpu
 
 two = torch.matmul(one, one).cuda()
-print('two device: ' + str(two.device))  # cpu
+
+
+def print_device():
+    print('zero device: ' + str(zero.device))
+    print('one device: ' + str(one.device))
+    print('two device: ' + str(two.device))
+
+
+print_device()
 
 
 @spaces.GPU
 def greet(n):
-    print('zero device: ' + str(zero.device))  # cuda
-    print('one device: ' + str(one.device))  # cuda
-    print('two device: ' + str(two.device))  # cuda
+    print('greeting')
+    print_device()
     return f"Hello {zero + n} Tensor"