How to create a tensor of size x based on a value (x) predicted by a network without breaking the computational graph

Hi All – Currently I’m trying to predict a kernel size through an auxiliary regression task which takes in features from a bottleneck of a CNN.

The regression task is predicting two values, kernel size and sigma - I am trying to generate a gaussian kernel based on that size and sigma - however, to generate the gaussian kernel, if I use torch.arange or torch.linspace it disconnects the predicted values from the computational graph. And if i still use a flag requires_grad on the torch.arange - it will be separate from the gradient from the predicted values - meaning i can’t backpropagate the gradients after computing the loss, therefore cannot train that portion of the network (I’ve visualized the graph through torchviz).

My question is - how can I generate a Gaussian kernel without having to disconnect the computation graph? How do i go about this problem? Any insight would be helpful! I’ve been at this problem for days and I can’t find a way out.

thank you!