Good day!
I am trying to train model to learn parameters for multiple Normal distributions, but receiving weird error.
Here is my code:
import torch
import torch.nn as nn
import random

def gen_data(num_classes):
    in_mean = [random.randint(-5, +5) for i in range(num_classes)]
    in_std = [random.random()+random.randint(1,3) for i in range(num_classes)]

    X = torch.tensor([])
    y = []
    idx = 0
    for mean, std in zip(in_mean, in_std):
        normal = torch.distributions.Normal(mean, std)
        n = random.randint(1, 100)
        sample = normal.sample((n, ),)
        X = torch.cat([X, sample])
        y += [idx]*n
        idx += 1
    y = torch.tensor(y).float()
    return X, y

class NaiveBayes(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.mean = nn.Parameter(torch.randn(num_classes), requires_grad=True)
        self.std = nn.Parameter(torch.abs(torch.randn(num_classes)), requires_grad=True)
        self.num_classes = num_classes

    def forward(self, x):
        x = x.repeat((self.num_classes, 1)).permute(1,0)
        normal = torch.distributions.Normal(self.mean, self.std)
        return normal.log_prob(x)

num_classes = 3
model = NaiveBayes(num_classes)
X, y = gen_data(num_classes)
criterion = torch.nn.MSELoss()

y_pred = model(X)          # shape: (num_samples, num_classes)
y_pred = y_pred.argmax(1)  # shape: (num_samples, )

loss = criterion(y, y_pred)
loss.backward()

After running this, the last line raises RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn. Are Normal distribution params trainable? If yes, what does this error suppose to mean?
torch==1.8.1
Thank you!