Hi @oheil, Thank you indeed it works whith your added correction although it’s lead to a slow down off the algorithm (it’s almost two times slower in distributed mode). This may be due to the need of allocating memory for each workers and it mean that I should probably look for other way to make my code faster, may be using the gpu.
Anyway thank you very much for your help on this !