forked from ccf-ai-infra/GPUCodeForces
29 lines
635 B
Python
29 lines
635 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, tau=0.5):
|
|
super().__init__()
|
|
self.tau = tau
|
|
|
|
def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
|
|
diff = y_true - y_pred
|
|
|
|
loss = torch.where(diff > 0, self.tau * diff, (self.tau - 1.0) * diff)
|
|
|
|
return loss.mean()
|
|
|
|
|
|
batch_size = 128
|
|
feature_dim = 512
|
|
|
|
|
|
def get_inputs():
|
|
y_pred = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
y_true = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
return [y_pred, y_true]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [0.5] |