import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH = 16
DEPTH = 64
HEIGHT = 64
WIDTH = 64

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, theta: torch.Tensor) -> torch.Tensor:

        return F.affine_grid(theta, size=(BATCH, 1, DEPTH, HEIGHT, WIDTH), align_corners=True)

def get_inputs():
 
    theta = torch.tensor([
        [1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0]
    ], dtype=torch.float32, device='cuda')
    
    theta = theta.unsqueeze(0).repeat(BATCH, 1, 1)
    return [theta]

def get_init_inputs():
    return []