Dimensionality of data and RV don't match in custom model

Describe the bug When fitting a model with custom likelihood, the dimensionality of the data and RV do not match.

HSSM version What is the version of HSSM you are using?

To Reproduce

import pandas as pd
import arviz as az
import numpy as np
import pandas as pd
import pytensor
import matplotlib.pyplot as plt 
import hssm
import ssms.basic_simulators

v_true_2 = 0.6
a_true_2 = 1.6 
z_true_2 = 0.6 
t_true_2 = 0.6 
shape_true = 5.0 
scale_true = 0.5 
c_true = 1.0 

test_dataset_2 = hssm.simulate_data(model='gamma_drift',
                                  theta=dict(v = v_true_2, 
                                             a = a_true_2,
                                             z = z_true_2,
                                             t = t_true_2, 
                                             shape = shape_true,
                                             scale = scale_true, 
                                             c = c_true  
                                             ),
                                             size = 500)

onnx_path_gamma = r"path_to_gamma_drift.onnx" 

gamma_drift_model = hssm.HSSM(data = test_dataset_2, 
                              model="gamma_drift",
                              model_config = {
                                  "list_params": ["v", "a", "z", "t", "shape", "scale", "c"],
                                  "bounds": { 
                                      "v": (-3.0, 3.0), 
                                      "a": (0.3, 3.0), 
                                      "z": (0.1, 0.9),
                                      "t": (0.001, 2.0),                                     
                                  }, 
                                  "backend": "jax",

                              },
                              shape = 5.0, 
                              scale = 0.5, 
                              c = 1.0, 
                              loglik = onnx_path_gamma,
                              loglik_kind = "approx_differentiable", 

                              )

Screenshots

Additional context This is not a problem with another custom model.

lnccbrown / HSSM

Dimensionality of data and RV don't match in custom model #511