import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.power = 16

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.linalg.matrix_power(x, self.power)

B = 1024 * 1024 
N = 4 

def get_inputs():
    x = torch.randint(0, 2, (B, N, N), device='cuda').float()
    return [x]

def get_init_inputs():
    return []