fsdp_review

Some eval and profile routines for fsdp

/mnist holds fsdp_mnist.py

This runs an fsdp CNN model for a default of 2 epochs to provide a review framework for FSDP.

Usage:

python fsdp_mnist.py  <optional flags> 

Flags:

-- use_offload = Activates cpu_offload
-- profile = Run PyTorch profile tracing
-- save-model = Save model (only working on nightlies)