File size: 379 Bytes
9b2bdf6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
import torch
from torch import nn
class WNormLoss(nn.Module):
def __init__(self, start_from_latent_avg=True):
super(WNormLoss, self).__init__()
self.start_from_latent_avg = start_from_latent_avg
def forward(self, latent, latent_avg=None):
if self.start_from_latent_avg:
latent = latent - latent_avg
return torch.sum(latent.norm(2, dim=(1, 2))) / latent.shape[0]
|