A lightweight wrapper for PyTorch that provides a simple declarative API for context switching between devices, distributed modes, mixed-precision, and PyTorch extensions.
borrows deepspeeds mpi detection for DDP which allows non deepspeed DDP to be easily used across multiple nodes when launched with openMPI (or the MPIOperator on k8s).
removes DataLoader batch size option as this can be directly pulled from the Stoke object.
borrows deepspeeds mpi detection for DDP which allows non deepspeed DDP to be easily used across multiple nodes when launched with openMPI (or the MPIOperator on k8s).
removes DataLoader batch size option as this can be directly pulled from the Stoke object.