import torch
import gradio as gr
import spaces
import nvdiffrast.torch as dr

@spaces.GPU
def pinfo():
    print("CUDA Version:", torch.version.cuda)
    print(torch.version.cuda)
    print(torch.cuda.is_available())
    zero = torch.Tensor([0]).cuda()
    print(zero.device)


def my_function(input_text):
    pinfo()
    return f"Received: {input_text}"


iface = gr.Interface(fn=my_function, inputs="text", outputs="text")
iface.launch()