Test the library with torchrun, as per @annahedstroem suggestion
Test the library with torchrun, as per @annahedstroem suggestion