Unable to use Pytorch with CUDA in Celery task
I have a celery task that uses torch library, which internally uses CUDA. When I run the task, it fails saying
&quot;Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method&quot;
When I browsed on this a little, I got this - https://github.com/celery/celery/issues/6036
This issue says Celery supports only fork and not spawn.
Is there any workaround/alternative to this?

you should load pytorch model and use it as global variable. The instance of model and task will run in the same process
from celery.signals import worker_process_init
# for more information about worker_process_init, read here:
# https://docs.celeryproject.org/en/stable/userguide/signals.html#worker-process-init

pytorch_model = None

@worker_process_init.connect()
def init_worker_process(**kwargs):
    &quot;&quot;&quot;
    load model before running tasks
    :param kwargs:
    :return:
    &quot;&quot;&quot;
    global pytorch_model
    pytorch_model = load_model()

@app.task
def predict_task(image: np.ndarray):
    return pytorch_model.predict(image)

