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 []
