forked from ccf-ai-infra/GPUCodeForces
31 lines
739 B
Python
31 lines
739 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, activation: str = 'relu') -> torch.Tensor:
|
|
|
|
fma_out = a * b + c
|
|
if activation == 'relu':
|
|
output = torch.relu(fma_out)
|
|
elif activation == 'sigmoid':
|
|
output = torch.sigmoid(fma_out)
|
|
else:
|
|
raise ValueError("Unsupported activation. Use 'relu' or 'sigmoid'")
|
|
return output
|
|
|
|
N = 2048 # Rows
|
|
M = 2048 # Columns
|
|
|
|
def get_inputs():
|
|
a = torch.randn(N, M)
|
|
b = torch.randn(N, M)
|
|
c = torch.randn(N, M)
|
|
return [a, b, c]
|
|
|
|
def get_init_inputs():
|
|
return []
|