import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        return torch.matmul(A, B)

def get_inputs():
    return [torch.randn(256, 512), torch.randn(512, 256)]

def get_init_inputs():
    return []