forked from ccf-ai-infra/GPUCodeForces
29 lines
644 B
Python
29 lines
644 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
|
|
margin = 1.0 - y_true * y_pred
|
|
|
|
loss_elementwise = F.relu(margin).pow(2)
|
|
|
|
return loss_elementwise.mean()
|
|
|
|
|
|
batch_size = 128
|
|
feature_dim = 512
|
|
|
|
|
|
def get_inputs():
|
|
y_pred = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
y_true = torch.randint(0, 2, (batch_size, feature_dim), dtype=torch.float32) * 2.0 - 1.0
|
|
return [y_pred, y_true]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |