Kernel crashes when running concordance_index()

I have a df with 40k rows and 21 variables. I am following the Churn prediction tutorial. csf_fit() works fine and takes 45min to run. But when I then run concordance_index() my session crashes and I lose my csf object.

I was able to reproduce the issue by running the example code for Conditional Survival Forest (CSF) but by increasing the N and number of features to:

# Generating N random samples 
N = 10000
dataset = sim.generate_data(num_samples = N, num_features=6)

I used the environment which the following Dockerfile provides:

FROM jupyter/scipy-notebook

RUN conda update -n base conda
RUN conda install pytorch-cpu torchvision-cpu -c pytorch
RUN conda install matplotlib pandas scikit-learn pyarrow progressbar scipy boost
RUN pip install --upgrade pip \
  && pip install pysurvival

square / pysurvival

Kernel crashes when running concordance_index() #6