画像スタイル転送テクノロジーにおけるスタイルの一貫性の問題

import torch
import torch.nn as nn
import torchvision.models as models

class StyleLoss(nn.Module):
    def __init__(self):
        super(StyleLoss, self).__init__()
        self.model = models.vgg19(pretrained=True).features[:23]
        self.layers = ['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1']
        
    def forward(self, input, target):
        input_features = self.model(input)
        target_features = self.model(target)
        
        loss = 0
        for layer in self.layers:
            input_style = self.gram_matrix(input_features[layer])
            target_style = self.gram_matrix(target_features[layer])
            loss += torch.mean(torch.square(input_style - target_style))
        
        return loss / len(self.layers)
        
    def gram_matrix(self, input):
        B, C, H, W = input.size()
        features = input.view(B * C, H * W)
        gram = torch.mm(features, features.t())
        
        return gram / (B * C * H * W)

class StyleTransferNet(nn.Module):
    def __init__(self, style_weight, content_weight, tv_weight):
        super(StyleTransferNet, self).__init__()
        self.style_loss = StyleLoss()
        self.content_loss = nn.MSELoss()
        self.tv_loss = nn.L1Loss()
        self.style_weight = style_weight
        self.content_weight = content_weight
        self.tv_weight = tv_weight
        
    def forward(self, input, target):
        style_loss = self.style_loss(input, target) * self.style_weight
        content_loss = self.content_loss(input, target) * self.content_weight
        tv_loss = self.tv_loss(input, target) * self.tv_weight
        
        return style_loss + content_loss + tv_loss