SVGP used to obtain attention: not optimizing

[fjsaezm]

Hello all!
I am new to torch! I am trying to develop a simple model in a simple example that:
Given an input $\mathbf{x} = (x_1,\cdots, x_n)$:

1. Uses a SVGP to predict the GP posterior $\mathbf{f} = SVGP(\mathbf{x})$
2. Takes samples from the marginal posterior and uses them as attention weigths for x : $w_i^{(s)} = sample(\mathbf{f})$
3. Using each sample, predicts a weighted representation of $\mathbf{x} \to \mathbf{r}^{(s)} = \sum_i w_i^{(s)} x_i$
4. Classifies $\mathbf{x}$ using $\mathbf{r}$.

The code I am using for the SVGP and sampling is the following:

# SVGP:
class GP(ApproximateGP):
    def __init__(self, inducing_points):
        
        # Store n_inducing
        self.num_Z = inducing_points.shape[0]
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(self.num_Z)    
        variational_strategy = VariationalStrategy(self,
                                                    inducing_points,
                                                    variational_distribution,
                                                    learn_inducing_locations=True)
        super(GP, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        self.covar_module.base_kernel.lengthscale = torch.tensor(1.0)
        
    def forward(self, x):
        mean = self.mean_module(x)
        covar = self.covar_module(x, diag = True)
        return gpytorch.distributions.MultivariateNormal(mean, torch.diag(covar))


# Attention computation:
class GPModel():
    def __init__(self, inducing_data,n_inducing = 32):
        super(GPModel, self).__init__(att_type = 'gp')
        
        batch_size = inducing_data.shape[0]
        bag_size = inducing_data.shape[1]
        inducing_data = inducing_data.reshape(batch_size * bag_size, *inducing_data.shape[2:])
        #print(inducing_data.shape)
        inducing_points = self.embedding(inducing_data)
        cluster_ids_x, cluster_centers = kmeans(
            X=inducing_points, num_clusters=n_inducing, distance='euclidean', device=torch.device('cuda:0')
        )
        self.mc_samples = 10
        self.attention_model = GP(inducing_points=cluster_centers)
        
    def attention(self,x):
        posterior = self.attention_model(x)
        #samples = posterior.get_base_samples(sample_shape = torch.Size([self.mc_samples]))
        samples = posterior.rsample(sample_shape = torch.Size([self.mc_samples]))
        return samples
        
    def kl_prior(self):
        return self.attention_model.variational_strategy.kl_divergence()

The model is currently not optimizing correctly, the full architecture works since I achieve good results using a deterministic linear model instead of a SVGP to compute the $w_i$, but the GP is not optimizing correctly to obtain correct attention weights, which worked previously in Tensorflow/GPFLOW. Maybe I am doing something wrong with the variables?

Thank you for your help :)