Luigi Task¶
- class aido.task.AIDOTask(*args, **kwargs)[source]
Bases:
TaskShallow wrapper around b2luigi.Task
- aido.task.torch_safe_wrapper(func: Callable, *args, **kwargs)[source]
b2luigi safe wrapper for calls to a torch function. Otherwise torch will raise ‘RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the ‘spawn’ start method’
We avoid this by first calling the function ‘func’ and excepting any errors. If that Error is raised by CUDA, we catch it and call that function again but inside a subprocess. If any further errors are raised afterwards, they are return. In case of no errors, we return the result of the function.