Error The actor ImplicitFunc is too large

Ray Libraries (Data, Train, Tune, Serve)
1

I’m using Ray version 2.9.0 and learning how to use Ray tune. Below is relevant code for tuning using How to use Tune with PyTorch — Ray 2.9.3 as reference example.

def train_moa(config):
   
    fold_num = 0

    train_idx = folds[folds['fold'] != fold_num].index
    val_idx = folds[folds['fold'] == fold_num].index
    train_folds = train.loc[train_idx].reset_index(drop=True)
    val_fold = train.loc[val_idx].reset_index(drop=True)
    train_target = y[train_idx]
    val_target = y[val_idx]

    for cat in cat_cols:
        train_folds[cat] = train_folds[cat].astype('category').cat.codes.values
        val_fold[cat] = val_fold[cat].astype('category').cat.codes.values
    cat_szs = [len(train_folds[col].astype('category').cat.categories) for col in cat_cols]
    emb_szs = [(size, min(50, (size+1)//2)) for size in cat_szs]

    train_dataset = TrainDataset(train_folds, cont_cols, cat_cols, train_target)
    val_dataset = TrainDataset(val_fold, cont_cols, cat_cols, val_target)

    train_loader = DataLoader(train_dataset, 
                              batch_size=CFG.batch_size, 
                              shuffle=True,
                              num_workers=4, 
                              pin_memory=True, 
                              drop_last=True
                              )
    val_loader = DataLoader(val_dataset, 
                            batch_size=CFG.batch_size, 
                            shuffle=False,
                            num_workers=4, 
                            pin_memory=True, 
                            drop_last=False
                            )
    
    # Model
    model = MOATabularModel(emb_szs, len(cont_cols), y.shape[1], CFG.hidden_units, p=CFG.dropout)
    model.to(device)

    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.SGD(
        model.parameters(), 
        lr=config["lr"], 
        momentum=config["momentum"]
        )
    
    for i in range(5):
        _ = train_func(train_loader, model, optimizer, criterion, device)
        val_loss, _ = val_func(val_loader, model, device) 

        train.report({"loss": val_loss})


search_space = {
    "lr": tune.sample_from(lambda spec: 10 ** (-10 * np.random.rand())),
    "momentum": tune.uniform(0.1, 0.9),
}

tuner = tune.Tuner(
    tune.with_parameters(train_moa),
    tune_config=tune.TuneConfig(
        num_samples=5,
    ),
    param_space=search_space,
)
results = tuner.fit()

Got error below. Any suggestions how to resolve?

File ~/.pyenv/versions/3.10.5/envs/kaggle/lib/python3.10/site-packages/ray/_private/utils.py:749, in check_oversized_function(pickled, name, obj_type, worker)
    738 error = (
    739     "The {} {} is too large ({} MiB > FUNCTION_SIZE_ERROR_THRESHOLD={}"
    740     " MiB). Check that its definition is not implicitly capturing a "
   (...)
    747     ray_constants.FUNCTION_SIZE_ERROR_THRESHOLD // (1024 * 1024),
    748 )
--> 749 raise ValueError(error)

ValueError: The actor ImplicitFunc is too large (392 MiB > FUNCTION_SIZE_ERROR_THRESHOLD=95 MiB). Check that its definition is not implicitly capturing a large array or other object in scope. Tip: use ray.put() to put large objects in the Ray object store.