| import gradio as gr | |
| import spaces | |
| import torch | |
| zero = torch.Tensor([0]).cuda() | |
| print(zero.device) # <-- 'cpu' 🤔 | |
| 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() | |