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__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.logdet(x)

B = 1024 * 128 
N = 8 

def get_inputs(): 
    Q = torch.randn(B, N, N, device='cuda', dtype=torch.float32)
    eye = torch.eye(N, device='cuda').unsqueeze(0)
    
    x = torch.matmul(Q, Q.transpose(1, 2)) + eye * 0.1
    
    return [x]

def get_init_inputs():
    return []