Can I optimize three parameters with a DistributionLambda layer?

I need to create a neural network layer to optimize the three parameters of a GeneralizedExtremeValue distribution (that are loc, scale, and concentration), but when more than one parameter is passed to a DistributionLambda layer, all training metrics are nan, and the output distribution is empty.

import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt
import tensorflow_probability as tfp
import tensorflow as tf
tfd = tfp.distributions
tfkl = tf.keras.layers
tfpl = tfp.layers

# INPUT DATASETS ==========================================================

# creating a GEV distribution 
dist = tfd.GeneralizedExtremeValue(loc=2, scale=1, concentration=0.1)
# sampling the GEV to create a dataset
dataset = dist.sample(10**(5))

# creating some noise on the sample
def add_eps (values):
    eps = np.random.randn(len(values))
    return values + eps

x_train = (0.75* add_eps(dataset[:(8*10**4)]))
y_train = dataset[:(8*10**4)]
x_test = (0.75* add_eps(dataset[(8*10**4):]))

# plotting the input datasets
fig=plt.figure()
sns.histplot(dataset, bins=100, 
             stat='probability', kde=True, 
             color='r', label='Sample from GEV: y_train')
sns.histplot(x_train, bins=100, 
             stat='probability',
             kde=True,
             color='b', label='Modified sample from GEV: x_train')
plt.legend()
plt.show()

Plotting the x_train and y_train, it seems that the input datasets are right. Anyway, the neural network outputs a distribution with nan instead of mean and stddev. I tried changing the batch_size, the epochs, and the number of outputs from the Dense layer, but nothing worked.

# CREATING BAYESIAN NETWORK ================================================

# model architecture
model = tf.keras.Sequential([
    tfkl.Dense(3, input_shape = (1,)),
    tfpl.DistributionLambda(
        lambda t: tfd.GeneralizedExtremeValue(loc=t[0], 
                             scale=t[1], concentration=t[2]))
    ])
negloglik = lambda y_true, y_pred: -y_pred.log_prob(y_true)
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), 
              loss=negloglik)

# Model Fitting
history = model.fit(x_train, y_train, 
                    validation_split=0.2, epochs=100, 
                    verbose=True, shuffle=True, 
                    batch_size =300)

# predictions
y_pred_mean = model(x_test.numpy().reshape(-1,1)).mean()
y_pred_std = model(x_test.numpy().reshape(-1,1)).stddev()