Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import triton | |
| import triton.language as tl | |
| def relu_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr): | |
| offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) | |
| mask = offs < n | |
| x = tl.load(x_ptr + offs, mask=mask) | |
| tl.store(y_ptr + offs, tl.maximum(x, 0.0), mask=mask) | |
| class ModelNew(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| def forward(self, x): | |
| x = x.contiguous() | |
| y = torch.empty_like(x) | |
| n = x.numel() | |
| relu_kernel[(triton.cdiv(n, 1024),)](x, y, n, BLOCK=1024) | |
| return y | |