import torch
import torch.nn as nn

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

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

    def forward(self, img1: torch.Tensor, img2: torch.Tensor, alpha: float) -> torch.Tensor:
     
        return img1 * alpha + img2 * (1.0 - alpha)

def get_inputs():

    img1 = torch.randint(low=0, high=255, size=(BATCH, CHANNELS, HEIGHT, WIDTH), device='cuda').float()
    img2 = torch.randint(low=0, high=255, size=(BATCH, CHANNELS, HEIGHT, WIDTH), device='cuda').float()
    

    return [img1, img2, 0.5]

def get_init_inputs():
    return []