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

BATCH = 64
CHANNELS = 64
HEIGHT = 128
WIDTH = 128
BLOCK_SIZE = 7 
KEEP_PROB = 0.9

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.block_size = BLOCK_SIZE
        self.keep_prob = KEEP_PROB
        self.gamma = None

    def calculate_gamma(self, x):
        return (1.0 - self.keep_prob) / (self.block_size ** 2) * \
               (x.shape[-2] * x.shape[-1]) / \
               ((x.shape[-2] - self.block_size + 1) * (x.shape[-1] - self.block_size + 1))

    def forward(self, x: torch.Tensor, rand_tensor: torch.Tensor) -> torch.Tensor:
        if self.gamma is None:
            self.gamma = self.calculate_gamma(x)

        mask = (rand_tensor < self.gamma).float()
        pad = self.block_size // 2
        mask = F.max_pool2d(mask, kernel_size=self.block_size, stride=1, padding=pad)
        mask = 1.0 - mask
        
        scale = 1.0 / self.keep_prob
        
        return x * mask * scale

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

def get_init_inputs():
    return []