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