File size: 650 Bytes
c14d03d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import torch
import torch.nn as nn
class IndentityWrapper(nn.Module):
def forward(self, loss: torch.Tensor) -> dict[str, torch.Tensor]:
return {"loss": loss}
class LossSumWrapper(nn.Module):
def __init__(self, weights: dict[str, float]):
super().__init__()
self.weights = weights
def forward(self,
loss_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
total_loss = 0
for loss_name, loss_val in loss_dict.items():
total_loss += loss_val * self.weights[loss_name]
output = {"loss": total_loss}
output.update(loss_dict)
return output
|