forked from ccf-ai-infra/GPUCodeForces
22 lines
537 B
Python
22 lines
537 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, in_features: int = 1024, out_features: int = 2048):
|
|
super().__init__()
|
|
self.linear = nn.Linear(in_features, out_features)
|
|
self.gelu = nn.GELU(approximate="none")
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.gelu(self.linear(x))
|
|
|
|
|
|
def get_init_inputs():
|
|
return {"in_features": 1024, "out_features": 2048}
|
|
|
|
|
|
def get_inputs():
|
|
B, T, D = 16, 512, 1024
|
|
x = torch.randn(B, T, D)
|
|
return x |