mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
@@ -0,0 +1,19 @@
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Tacotron2Loss(nn.Module):
|
||||
def __init__(self):
|
||||
super(Tacotron2Loss, self).__init__()
|
||||
|
||||
def forward(self, model_output, targets):
|
||||
mel_target, gate_target = targets[0], targets[1]
|
||||
mel_target.requires_grad = False
|
||||
gate_target.requires_grad = False
|
||||
gate_target = gate_target.view(-1, 1)
|
||||
|
||||
mel_out, mel_out_postnet, gate_out, _ = model_output
|
||||
gate_out = gate_out.view(-1, 1)
|
||||
mel_loss = nn.MSELoss()(mel_out, mel_target) + \
|
||||
nn.MSELoss()(mel_out_postnet, mel_target)
|
||||
gate_loss = nn.BCEWithLogitsLoss()(gate_out, gate_target)
|
||||
return mel_loss + gate_loss
|
||||
Reference in New Issue
Block a user