import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, shift):
        super(Model, self).__init__()
        self.shift = shift

    def forward(self, x):
        return torch.sigmoid(torch.logit(x)) + self.shift

batch_size = 4096
dim = 1024

def get_inputs():
    x = torch.rand(batch_size, dim) * 0.999 + 0.0005
    return [x]

def get_init_inputs():
    return [0.5]