# bcewithlogitsloss_torch.py
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 4096
FEATURE_DIM = 256

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.criterion = nn.BCEWithLogitsLoss(reduction='mean') 

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
        return self.criterion(logits, target)

def get_inputs():

    logits = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
    target = torch.randint(0, 2, (BATCH_SIZE, FEATURE_DIM), dtype=torch.float32)
    return [logits, target]

def get_init_inputs():
    return []