import gradio as gr
import spaces
import torch

zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' 🤔


@spaces.GPU
def greet(n):
    print(zero.device) # <-- 'cuda:0' 🤗
    if torch.cuda.is_available():
        total_memory_bytes = torch.cuda.get_device_properties(0).total_memory
        gb =  total_memory_bytes / (1024**3)
        print(f"VRAM 0: {gb:.2f} GB")
    else:
        print("CUDA is NOT available")
    return f"Hello {zero + n} Tensor"

demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text(), share=True)
demo.launch()