forked from ccf-ai-infra/GPUCodeForces
33 lines
716 B
Python
33 lines
716 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, tl=-1.0, al=0.1, tr=1.0, ar=0.1):
|
|
super().__init__()
|
|
self.tl = tl
|
|
self.al = al
|
|
self.tr = tr
|
|
self.ar = ar
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
y_left = self.tl + self.al * (x - self.tl)
|
|
|
|
y_right = self.tr + self.ar * (x - self.tr)
|
|
|
|
y_mid_and_right = torch.where(x < self.tr, x, y_right)
|
|
|
|
return torch.where(x <= self.tl, y_left, y_mid_and_right)
|
|
|
|
|
|
batch_size = 128
|
|
feature_dim = 512
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
return [x]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [-1.0, 0.1, 1.0, 0.1] |