torch tensor object error while using bayesian for market mix modeling

Question

I am developing a MMM project which uses the below adstock transformation function using torch:

def adstock_geometric_pytensor(x, theta):
    x = torch.tensor(x,dtype=torch.float32)
    theta = torch.tensor(theta, dtype=torch.float32)
    
    def adstock_geometric_recurrence_pytensor(index, input_x, decay_x, theta):
        decay_x[index] = input_x + theta * decay_x[index-1]
        return decay_x

    len_observed = x.shape[0]

    x_decayed = torch.zeros_like(x)
    x_decayed[0] = x[0]

    for index in range(1,len_observed):
        x_decayed[index] = adstock_geometric_recurrence_pytensor(index,x[index],x_decayed,theta)
    
    return x_decayed

This is futher used down below while defining the model priors:

%%time

transform_variables = ["trend", "season", "holiday", 'Audio_s','Cable_s']
    
delay_channels = ['Audio_s','Cable_s']
media_channels = ['Audio_s','Cable_s']
control_variables = ["trend", "season", "holiday"]

target = "Joins"

data_transformed = final_data.copy()

numerical_encoder_dict = {}

for feature in transform_variables:
    scaler = MinMaxScaler()
    original = final_data[feature].values.reshape(-1, 1)
    transformed = scaler.fit_transform(original)
    data_transformed[feature] = transformed
    numerical_encoder_dict[feature] = scaler

dependent_transformation = None
original = final_data[target].values
data_transformed[target] = original 

# #Model 1:
response_mean = []
with pm.Model() as model_2:
    for channel_name in delay_channels:
        
        print(f"Delay Channels: Adding {channel_name}")
        
        x = data_transformed[channel_name].values
        
        adstock_param = pm.Beta(f"{channel_name}_adstock",3,3)
        pidata = pm.sample(1000,random_seed=43)
        adstock_param = np.average(pidata.posterior[f"{channel_name}_adstock"].to_numpy().flatten())
        
        saturation_gamma = pm.Beta(f"{channel_name}_gamma",2,2)
        pidata = pm.sample(1000,random_seed=43)
        saturation_gamma = np.average(pidata.posterior[f"{channel_name}_gamma"].to_numpy().flatten())
        
        saturation_alpha = pm.Beta(f"{channel_name}_alpha",3,1)
        pidata = pm.sample(1000,random_seed=43)
        saturation_alpha = np.average(pidata.posterior[f"{channel_name}_alpha"].to_numpy().flatten())
        
        
        x_new = adstock_geometric_pytensor(x, adstock_param)
        x_new_sliced = x_new[START_ANALYSIS_INDEX:END_ANALYSIS_INDEX]
        saturation_tensor = saturation_hill_pymc(x_new_sliced, saturation_alpha, saturation_gamma)
    

However this gives me the below error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File <timed exec>:51

Cell In[67], line 47, in adstock_geometric_pytensor(x, theta)
     44 x_decayed[0] = x[0]
     46 for index in range(1,len_observed):
---> 47     x_decayed[index] = adstock_geometric_recurrence_pytensor(index,x[index],x_decayed,theta)
     49 return x_decayed

RuntimeError: expand(torch.FloatTensor{[365]}, size=[]): the number of sizes provided (0) must be greater or equal to the number of dimensions in the tensor (1)

Not sure what is causing the error, is it the function adstock_geometric_pytensor or the way I am defining the priors under pm.Model()

If someone has experience into market mix modeling with adstock transformations using pymc latest version which doesn't use Theano anymore, can he/she kindly guide me here.