sofhiaazzhr's picture
Add GPU + API smoke test
e9d7853 verified
Raw History Blame Contribute Delete
750 Bytes
import gradio as gr
import spaces
import torch
@spaces.GPU
def check_gpu():
"""Runs on ZeroGPU — proves the GPU is actually allocated to this call."""
available = torch.cuda.is_available()
name = torch.cuda.get_device_name(0) if available else "none"
x = torch.tensor([1.0]).cuda() # forces a real GPU op
return {
"cuda_available": available,
"device": str(x.device), # should be 'cuda:0'
"gpu_name": name,
}
with gr.Blocks() as demo:
gr.Markdown("# knowledge-pipeline — GPU + API smoke test")
out = gr.JSON(label="result")
gr.Button("Check GPU").click(fn=check_gpu, inputs=None, outputs=out,
api_name="check_gpu")
demo.launch()