File size: 1,077 Bytes
d4f211f c311b2d d4f211f 563fb2c d4f211f 376c53e 0237ca0 c311b2d 0237ca0 c311b2d 5326ed4 f384f81 117e6d7 563fb2c e89c6d2 c311b2d cdab2e6 c311b2d e89c6d2 c311b2d e89c6d2 117e6d7 c311b2d 117e6d7 c311b2d 117e6d7 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 53 54 55 |
import gradio as gr
import spaces
import torch
import os
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(n: int):
print(f'\n===process step {n}')
print('cuda visible devices: ' + os.getenv('CUDA_VISIBLE_DEVICES'))
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(2)
return f"Hello {n} Tensor"
def func(n):
# step1, on cpu
process(1)
# step2, on gpu
res = greet(n)
# step3, on cpu
process(3)
return res
gr.Interface(fn=func, inputs=gr.Number(), outputs=gr.Text()).launch()
|