text: #
  # SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  # SPDX-License-Identifier: Apache-2.0
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  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
  