I don't fully understand why but for me this solved the issue:

if __name__ == '__main__':
    import torch
    torch.multiprocessing.set_start_method('spawn')
    main()