File size: 565 Bytes
06d7fc5 965d9e4 06d7fc5 965d9e4 5002466 965d9e4 06d7fc5 965d9e4 06d7fc5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | 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()
|