Databricks notebook cell gets hung up but code completes

Question

Have been running into an issue when running a pymc-marketing model in a Databricks notebook. The cell that fits the model gets hung up and the progress bar stops moving, however the code completes and dumps all needed output into a folder. After the code completes I have to then detach the notebook since hitting Interrupt doesn't respond. I took a peek at the cluster logs and can confirm everything runs as expected (see screenshot!).

Any ideas the issue here or have you run into the same issue??

Here's some context:

Data

• sample size: 22 weeks of data - small size here just for easy of model runs for now.
• input vars: 11

Model Specification

{
"meida_channel": [
    "x1",
    "x2",
    "x3",
    "x4",
    "x5",
    "x6",
    "x7"
],
"control": [
    "c1",
    "c2",
    "c3",
    "c4"
],
"response_var": "units",
"date_var": "end_date",
"beta_channel": {
    "dist": "TruncatedNormal",
    "kwargs":{
        "mu": [stuff here],
        "sigma": 0.001,
        "lower": 0
    }
},
"alpha": {
    "dist": "Beta",
    "kwargs": {
        "alpha": 1,
        "beta": 3
    }
},
"lam": {
    "dist": "Gamma",
    "kwargs": {
        "alpha": 2,
        "beta": 1
    }
},
"gamma_control": {
    "dist": "Normal",
    "kwargs": {
        "mu": [stuff here],
        "sigma": 0.001
    }
}}

sampler_kwargs = {
    "draws": 8_000,
    "target_accept": 0.95,
    "chains": 4,
    "random_seed": rng,
    # "nuts": {"max_treedepth": 15},
}