import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_boundaries=128):
        super().__init__()
        self.boundaries = nn.Parameter(
            torch.linspace(0, 100, steps=num_boundaries), 
            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] 