forked from ccf-ai-infra/GPUCodeForces
38 lines
759 B
Python
38 lines
759 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
original_shape = x.shape
|
|
x_flat = x.flatten()
|
|
D = x_flat.size(0)
|
|
|
|
indices = torch.arange(D, device=x.device)
|
|
|
|
mask_even = (indices % 2 == 0)
|
|
|
|
result = torch.zeros_like(x_flat)
|
|
|
|
result[mask_even] = F.relu(x_flat[mask_even])
|
|
|
|
result[~mask_even] = -F.relu(-x_flat[~mask_even])
|
|
|
|
return result.view(original_shape)
|
|
|
|
|
|
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 [] |