ахуеете, работает так же быстро как у торча (а иногда даже быстрее), однако код понятный
pip install tiny-torch-nevmenkoimport torch
import tiny_torch.functional as F
x = torch.tensor([-1.0, 0.0, 1.0], device="cuda")
print(F.relu(x))
# tensor([0., 0., 1.], device='cuda:0')
Here I'll reproduce some torch functions to have a sdome experience with torch_cpp_extentions and CUDA