mirror of
https://github.com/huggingface/kernels.git
synced 2025-10-20 21:10:02 +08:00
373 B
373 B
kernels
import torch
from kernels import get_kernel
# Download optimized kernels from the Hugging Face hub
layer_norm_kernels = get_kernel("kernels-community/layer-norm")
# Initialize torch Module
optimized_layer_norm_layer = layer_norm_kernels.DropoutAddLayerNorm(128).cuda()
# Forward
x = torch.randn(128).cuda()
print(optimized_layer_norm_layer(x))