forked from ccf-ai-infra/GPUCodeForces
20 lines
398 B
Python
20 lines
398 B
Python
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] |