import torch
import torch.nn as nn
import torchvision.transforms.functional as F


BATCH = 64
CHANNELS = 3
HEIGHT = 1024
WIDTH = 1024

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

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

        return F.adjust_gamma(x, self.gamma)

def get_inputs():

    x = torch.rand(BATCH, CHANNELS, HEIGHT, WIDTH, device='cuda', dtype=torch.float32)
    return [x]

def get_init_inputs():
    return []