forked from ccf-ai-infra/GPUCodeForces
26 lines
573 B
Python
26 lines
573 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, num_bins=128):
|
|
super().__init__()
|
|
|
|
self.boundaries = nn.Parameter(
|
|
torch.linspace(0, 100, steps=num_bins),
|
|
requires_grad=False
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
|
return torch.bucketize(x, self.boundaries, right=False)
|
|
|
|
num_elements = 100 * 1024 * 1024
|
|
shape = (num_elements, )
|
|
|
|
def get_inputs():
|
|
|
|
x = torch.rand(shape, dtype=torch.float32) * 100.0
|
|
return [x]
|
|
|
|
def get_init_inputs():
|
|
return [128] |