forked from ccf-ai-infra/GPUCodeForces
32 lines
646 B
Python
32 lines
646 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, min_val, max_val, eps=1e-5):
|
|
super(Model, self).__init__()
|
|
self.min_val = min_val
|
|
self.max_val = max_val
|
|
self.eps = eps
|
|
|
|
def forward(self, x):
|
|
mean = x.mean(dim=-1, keepdim=True)
|
|
std = x.std(dim=-1, keepdim=True)
|
|
|
|
norm = (x - mean) / (std + self.eps)
|
|
return torch.clamp(norm, self.min_val, self.max_val)
|
|
|
|
|
|
batch_size = 16
|
|
dim = 256
|
|
min_val = -1.0
|
|
max_val = 1.0
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, dim) * 10.0
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [min_val, max_val] |