File size: 950 Bytes
d4f211f c311b2d d4f211f 376c53e 0237ca0 c311b2d 0237ca0 c311b2d 5326ed4 f384f81 c311b2d e89c6d2 c311b2d cdab2e6 c311b2d e89c6d2 c311b2d e89c6d2 c311b2d d4f211f 1ee0472 c311b2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import gradio as gr
import spaces
import torch
torch.set_default_device('cuda')
class LinearModel(torch.nn.Module):
def __init__(self):
super(LinearModel, self).__init__()
self.linear = torch.nn.Linear(1, 1, device='cuda')
def forward(self, x):
out = self.linear(x)
print('weight device: ' + str(self.linear.weight.device))
return out
def process():
print('cuda avaliable: ' + str(torch.cuda.is_available()))
print('cuda device count: ' + str(torch.cuda.device_count()))
model = LinearModel().cuda()
x = torch.ones(1, device='cuda')
y = model(x)
print(model)
print(x)
print(y)
@spaces.GPU
def greet(n):
process()
return f"Hello {n} Tensor"
def func(n):
# step1, on cpu
process()
# step2, on gpu
res = greet(n)
# step3, on cpu
process()
return res
gr.Interface(fn=func, inputs=gr.Number(), outputs=gr.Text()).launch()
|