I want to make a custom loss function of MSE Loss by doing GMM computation for the result.
The MSE Loss will be computed using the GMM result and target value.
Here is my code

class StyleLoss(nn.Module):
def __init__(self, target_feature, ncomp, initial_mus, initial_covs, initial_priors):

    super(StyleLoss, self).__init__()
    self.gmm = GMM(ncomp, initial_mus, initial_covs, initial_priors)
    self.target, self.log_likelihood_target = self.gmm.inference(target_feature)
    self.target = torch.tensor(self.target, dtype=torch.float).to(device)

def forward(self, ctx, input_style):

    input_style = input_style.detach().cpu().numpy()
    image_batch, image_channels, image_height, image_width = input_style.shape
    input_style_reshaped = np.reshape(input_style, (-1, image_channels))

    self.input_gmm, self.log_likelihood_input = self.gmm.inference(input_style_reshaped)
    self.input_gmm = torch.tensor(self.input_gmm, dtype=torch.float).to(device)
    self.loss = F.mse_loss(self.input_gmm, self.target)

    self.loss = torch.tensor(self.loss, requires_grad=True, dtype=torch.float).to(device)
    input_style = torch.tensor(input_style, dtype=torch.float).to(device)

    return input_style


I have read many discussion forum posts but still doesn’t get the proper method to make this custom loss function, because when I do it this, the gradient will not be calculated. Hope everyone can help me. Bunch of thanks before !