FeatureMatching Loss feat_scale not used?

[ZhikangNiu]

audiocraft/audiocraft/adversarial/losses.py

Lines 217 to 223 in adf0b04

	feat_scale = torch.tensor(0., device=fmap_fake[0].device)
	n_fmaps = 0
	for (feat_fake, feat_real) in zip(fmap_fake, fmap_real):
	assert feat_fake.shape == feat_real.shape
	n_fmaps += 1
	feat_loss += self.loss(feat_fake, feat_real)
	feat_scale += torch.mean(torch.abs(feat_real))

The feat_scale isn't used in training. Maybe the correct is:

def forward(self, fmap_fake: tp.List[torch.Tensor], fmap_real: tp.List[torch.Tensor]) -> torch.Tensor:
        assert len(fmap_fake) == len(fmap_real) and len(fmap_fake) > 0
        feat_loss = torch.tensor(0., device=fmap_fake[0].device)
        feat_scale = torch.tensor(0., device=fmap_fake[0].device)
        n_fmaps = 0
        for (feat_fake, feat_real) in zip(fmap_fake, fmap_real):
            assert feat_fake.shape == feat_real.shape
            n_fmaps += 1
            feat_loss += self.loss(feat_fake, feat_real)
            feat_scale += torch.mean(torch.abs(feat_real))
            feat_loss /= feat_scale
        if self.normalize:
            feat_loss /= n_fmaps

        return feat_loss