forked from ccf-ai-infra/GPUCodeForces
24 lines
482 B
Python
24 lines
482 B
Python
# l1loss_torch.py
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
BATCH_SIZE = 16
|
|
DIM = 16384 * 16
|
|
|
|
class Model(nn.Module):
|
|
|
|
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
|
return F.l1_loss(pred, target, reduction='mean')
|
|
|
|
def get_inputs():
|
|
|
|
pred = torch.randn(BATCH_SIZE, DIM, dtype=torch.float32)
|
|
target = pred + torch.rand_like(pred) * 0.1
|
|
return [pred, target]
|
|
|
|
def get_init_inputs():
|
|
|
|
return [] |