import torch
import torch.nn as nn


N, C, H, W = 32, 64, 128, 128

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

    def forward(self, a, b):

        a_broad, b_broad = torch.broadcast_tensors(a, b)

        return torch.stack([a_broad, b_broad], dim=0).contiguous()

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

    def forward(self, a, b):
        return self.op(a, b)

def get_inputs():
    a = torch.randn(N, 1, H, W, dtype=torch.float32)
    b = torch.randn(1, C, H, W, dtype=torch.float32)
    return [a, b]

def get_init_inputs():
    return []