File size: 312 Bytes
7f6b5fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import gradio as gr
import spaces
import torch

zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' πŸ€”

@spaces.GPU
def chat(prompt):
    print(zero.device) # <-- 'cuda:0' πŸ€—
    return f"Hello {zero + n} Tensor"

gr.Interface(
    fn=chat, 
    inputs=gr.Text(), 
    outputs=gr.Text()
).launch()