Seeking assistance with ValueError: Model 'linear_regression' not found for {symbol}.

[EuclideanWild]

I have a data frame with the columns date, open, high, low, close, brecken, lofi, symbol, and pred already created. I registered it with
pybroker.register_columns('brecken', 'lofi', 'prev_pct_change', 'pred')

Then I defined the model:

# Define the custom model training function

def train_model(symbol, train_data, test_data):
    X_train = train_data[['close', 'brecken', 'lofi']]
    y_train = train_data['pred']
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    X_test = test_data[['close', 'brecken', 'lofi']]
    y_test = test_data['pred']
    y_pred = model.predict(X_test)

    r2 = r2_score(y_test, y_pred)
    print(f'{symbol} R^2={r2}')
    
    return model

# Register the model
linear_regression = pybroker.model('linear_regression', train_model)

Then I defined the execution logic:

# Define the execution logic
def strategy_logic(ctx):
    if not ctx.long_pos() and ctx.preds('linear_regression')[-1] > 0:
        ctx.buy_shares = 100
    elif ctx.long_pos() and ctx.preds('linear_regression')[-1] < 0:
        ctx.sell_shares = 100

And then I tried to backtest:

strategy = Strategy(data, '12/20/2023', '05/13/2024')
strategy.add_execution(strategy_logic, ['AAPL'], models=linear_regression)
result = strategy.backtest()
result.orders

However, when I run it I get the error:

ValueError: Model 'linear_regression' not found for AAPL.

Does anyone know how I can get this to work?