Utilisation de PyTorch avec du céleri

J'essaye d'exécuter un modèle PyTorch dans une application Django. Comme il n'est pas recommandé d'exécuter les modèles (ou toute tâche de longue durée) dans les vues, j'ai décidé de l'exécuter dans une tâche Celery. Mon modèle est assez grand et il faut environ 12 secondes pour charger et environ 3 secondes pour en déduire. C'est pourquoi j'ai décidé que je ne pouvais pas me permettre de le charger à chaque demande. J'ai donc essayé de le charger dans les paramètres et de l'enregistrer là-bas pour que l'application l'utilise. Donc mon dernier schéma est:

• Lorsque l'application Django démarre, dans les paramètres, le modèle PyTorch est chargé et accessible depuis l'application.
• Lorsque views.py reçoit une demande, cela retarde une tâche de céleri
• La tâche céleri utilise le settings.model pour déduire le résultat

Le problème ici est que la tâche de céleri génère l'erreur suivante lors de la tentative d'utilisation du modèle

@task
def getResult(name):
    print("Executing on GPU:", torch.cuda.is_available())
    if os.path.isfile(name):
        try:
            outpath = model_inference(name)
            os.remove(name)
            return outpath
        except OSError as e:
            print("Error", name, "doesn't exist")
    return ""

Voici le code dans mon settings.py chargeant le modèle:

if sys.argv and sys.argv[0].endswith('celery') and 'worker' in sys.argv: #In order to load only for the celery worker
    import torch
    torch.cuda.init()
    torch.backends.cudnn.benchmark = True
    load_model_file()

Et le code de tâche

[2020-08-29 09:03:04,015: ERROR/ForkPoolWorker-1] Task app.tasks.task[458934d4-ea03-4bc9-8dcd-77e4c3a9caec] raised unexpected: RuntimeError("Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method")
Traceback (most recent call last):
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/celery/app/trace.py", line 412, in trace_task
    R = retval = fun(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/celery/app/trace.py", line 704, in __protected_call__
    return self.run(*args, **kwargs)
  /*...*/
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/torch/cuda/__init__.py", line 191, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. " + msg)
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

L'impression dans la tâche indique "Executing on GPU: true"

J'ai essayé de définir torch.multiprocessing.set_start_method('spawn') dans le settings.py avant et après torch.cuda.init() mais cela donne la même erreur.